fix(pypi): rewrite RECORD file entries for extracted .data contents

news/4025.fixed.md

@@ -0,0 +1,3 @@
+(pypi) Fixed {obj}`RECORD` file paths for extracted `.data` directory contents
+so that {obj}`importlib.metadata.files()` correctly locates installed
+distribution files ([#4025](https://github.com/bazel-contrib/rules_python/pull/4025)).

python/private/pypi/BUILD.bazel

@@ -41,6 +41,15 @@ alias(
     visibility = ["//visibility:public"],
 )
 
+alias(
+    name = "wheel_record_rewriter",
+    actual = select({
+        "@platforms//os:windows": "wheel_record_rewriter.ps1",
+        "//conditions:default": "wheel_record_rewriter.sh",
+    }),
+    visibility = ["//visibility:public"],
+)
+
 exports_files(
     srcs = ["deps.bzl"],
     visibility = ["//tools/private/update_deps:__pkg__"],
@@ -490,11 +499,20 @@ bzl_library(
     ],
 )
 
+bzl_library(
+    name = "gen_wheel_record",
+    srcs = ["gen_wheel_record.bzl"],
+    deps = [
+        "//python/private:common",
+    ],
+)
+
 bzl_library(
     name = "whl_library_targets",
     srcs = ["whl_library_targets.bzl"],
     deps = [
         ":env_marker_setting",
+        ":gen_wheel_record",
         ":labels",
         ":namespace_pkgs",
         ":pep508_deps",

python/private/pypi/gen_wheel_record.bzl

@@ -0,0 +1,83 @@
+"""Rule for generating platform-specific RECORD files."""
+
+load("//python/private:common.bzl", "is_windows_platform")
+
+def _gen_wheel_record_impl(ctx):
+    is_windows = is_windows_platform(ctx)
+    rewriter_file = ctx.files._wheel_record_rewriter[0]
+    out_files = []
+
+    for in_file in ctx.files.srcs:
+        dist_info_name = in_file.dirname.rpartition("/")[2]
+        if dist_info_name:
+            if dist_info_name.endswith(".dist-info"):
+                data_dir_basename = (
+                    dist_info_name[:-len(".dist-info")] + ".data"
+                )
+            else:
+                data_dir_basename = dist_info_name + ".data"
+            out_file = ctx.actions.declare_file(
+                "site-packages/{}/RECORD".format(dist_info_name),
+            )
+        else:
+            data_dir_basename = "data"
+            out_file = ctx.actions.declare_file("site-packages/RECORD")
+
+        out_files.append(out_file)
+
+        action_args = ctx.actions.args()
+        inputs = depset([in_file, rewriter_file])
+
+        if rewriter_file.path.endswith(".ps1"):
+            action_exe = "powershell.exe"
+            action_args.add_all([
+                "-ExecutionPolicy",
+                "Bypass",
+                "-NoProfile",
+                "-File",
+                rewriter_file,
+            ])
+        else:
+            action_exe = (
+                ctx.attr._wheel_record_rewriter[DefaultInfo].files_to_run
+            )
+
+        action_args.add(in_file)
+        action_args.add(out_file)
+        action_args.add("windows" if is_windows else "unix")
+        action_args.add(data_dir_basename)
+
+        ctx.actions.run(
+            inputs = inputs,
+            outputs = [out_file],
+            executable = action_exe,
+            arguments = [action_args],
+            mnemonic = "PyGenWheelRecord",
+            progress_message = "Generating wheel RECORD %{output}",
+            toolchain = None,
+        )
+
+    return [
+        DefaultInfo(files = depset(out_files)),
+    ]
+
+gen_wheel_record = rule(
+    implementation = _gen_wheel_record_impl,
+    attrs = {
+        "srcs": attr.label_list(
+            doc = "The original RECORD files to rewrite.",
+            mandatory = True,
+            allow_files = True,
+        ),
+        "_wheel_record_rewriter": attr.label(
+            default = "//python/private/pypi:wheel_record_rewriter",
+            allow_files = True,
+            cfg = "exec",
+        ),
+        "_windows_constraints": attr.label_list(
+            default = [
+                "@platforms//os:windows",
+            ],
+        ),
+    },
+)

python/private/pypi/wheel_record_rewriter.ps1

@@ -0,0 +1,66 @@
+[CmdletBinding()]
+param(
+    [Parameter(Position=0, Mandatory=$true)]
+    [string]$InFile,
+
+    [Parameter(Position=1, Mandatory=$true)]
+    [string]$OutFile,
+
+    [Parameter(Position=2, Mandatory=$true)]
+    [string]$TargetOs,
+
+    [Parameter(Position=3, Mandatory=$true)]
+    [string]$DataDirBasename
+)
+
+$ErrorActionPreference = "Stop"
+
+$dataPrefix = "$DataDirBasename/"
+$quotedDataPrefix = "`"$DataDirBasename/"
+
+if ($TargetOs -eq "windows") {
+    $dataRepl = "../../"
+    $headersRepl = "../../Include/"
+    $platlibRepl = ""
+    $purelibRepl = ""
+    $scriptsRepl = "../../Scripts/"
+} else {
+    $dataRepl = "../../../"
+    $headersRepl = "../../../include/"
+    $platlibRepl = ""
+    $purelibRepl = ""
+    $scriptsRepl = "../../../bin/"
+}
+
+$lines = Get-Content -Path $InFile
+$outLines = [System.Collections.Generic.List[string]]::new()
+$Utf8NoBom = New-Object System.Text.UTF8Encoding $False
+
+foreach ($line in $lines) {
+    if ($line.StartsWith($quotedDataPrefix)) {
+        $quote = "`""
+        $rest = $line.Substring($quotedDataPrefix.Length)
+    } elseif ($line.StartsWith($dataPrefix)) {
+        $quote = ""
+        $rest = $line.Substring($dataPrefix.Length)
+    } else {
+        $outLines.Add($line)
+        continue
+    }
+
+    if ($rest.StartsWith("purelib/")) {
+        $outLines.Add($quote + $purelibRepl + $rest.Substring(8))
+    } elseif ($rest.StartsWith("platlib/")) {
+        $outLines.Add($quote + $platlibRepl + $rest.Substring(8))
+    } elseif ($rest.StartsWith("scripts/")) {
+        $outLines.Add($quote + $scriptsRepl + $rest.Substring(8))
+    } elseif ($rest.StartsWith("headers/")) {
+        $outLines.Add($quote + $headersRepl + $rest.Substring(8))
+    } elseif ($rest.StartsWith("data/")) {
+        $outLines.Add($quote + $dataRepl + $rest.Substring(5))
+    } else {
+        $outLines.Add($line)
+    }
+}
+
+[System.IO.File]::WriteAllText($OutFile, ($outLines -join "`n") + "`n", $Utf8NoBom)

python/private/pypi/wheel_record_rewriter.sh

@@ -0,0 +1,60 @@
+#!/bin/sh
+set -eu
+
+IN="$1"
+OUT="$2"
+TARGET_OS="$3"
+DATA_DIR_BASENAME="$4"
+
+DATA_PREFIX="${DATA_DIR_BASENAME}/"
+QUOTED_DATA_PREFIX="\"${DATA_DIR_BASENAME}/"
+
+if [ "$TARGET_OS" = "windows" ]; then
+  DATA_REPL="../../"
+  HEADERS_REPL="../../Include/"
+  PLATLIB_REPL=""
+  PURELIB_REPL=""
+  SCRIPTS_REPL="../../Scripts/"
+else
+  DATA_REPL="../../../"
+  HEADERS_REPL="../../../include/"
+  PLATLIB_REPL=""
+  PURELIB_REPL=""
+  SCRIPTS_REPL="../../../bin/"
+fi
+
+awk -v data_prefix="$DATA_PREFIX" \
+    -v quoted_data_prefix="$QUOTED_DATA_PREFIX" \
+    -v data_repl="$DATA_REPL" \
+    -v headers_repl="$HEADERS_REPL" \
+    -v platlib_repl="$PLATLIB_REPL" \
+    -v purelib_repl="$PURELIB_REPL" \
+    -v scripts_repl="$SCRIPTS_REPL" '
+{
+  line = $0
+  quote = ""
+  if (substr(line, 1, length(quoted_data_prefix)) == quoted_data_prefix) {
+    quote = "\""
+    rest = substr(line, length(quoted_data_prefix) + 1)
+  } else if (substr(line, 1, length(data_prefix)) == data_prefix) {
+    rest = substr(line, length(data_prefix) + 1)
+  } else {
+    print line
+    next
+  }
+
+  if (substr(rest, 1, 8) == "purelib/") {
+    print quote purelib_repl substr(rest, 9)
+  } else if (substr(rest, 1, 8) == "platlib/") {
+    print quote platlib_repl substr(rest, 9)
+  } else if (substr(rest, 1, 8) == "scripts/") {
+    print quote scripts_repl substr(rest, 9)
+  } else if (substr(rest, 1, 8) == "headers/") {
+    print quote headers_repl substr(rest, 9)
+  } else if (substr(rest, 1, 5) == "data/") {
+    print quote data_repl substr(rest, 6)
+  } else {
+    print line
+  }
+}
+' "$IN" > "$OUT"

python/private/pypi/whl_extract.bzl

@@ -36,20 +36,7 @@ def whl_extract(rctx, *, whl_path, logger):
     # Get the <prefix>.dist_info dir name
     data_dir = dist_info_dir.dirname.get_child(dist_info_dir.basename[:-len(".dist-info")] + ".data")
     if data_dir.exists:
-        for prefix, dest_prefix in {
-            # https://docs.python.org/3/library/sysconfig.html#posix-prefix
-            # We are taking this from the legacy whl installer config
-            "data": "data",
-            "headers": "include",
-            # In theory there may be directory collisions here, so it would be best to
-            # merge the paths here. We are doing for quite a few levels deep. What is
-            # more, this code has to be reasonably efficient because some packages like
-            # to not put everything to the top level, but to indicate explicitly if
-            # something is in `platlib` or `purelib` (e.g. libclang wheel).
-            "platlib": "site-packages",
-            "purelib": "site-packages",
-            "scripts": "bin",
-        }.items():
+        for prefix, dest_prefix in _DATA_CATEGORIES.items():
             src = data_dir.get_child(prefix)
             if not src.exists:
                 # The prefix does not exist in the wheel, we can continue
@@ -61,9 +48,32 @@ def whl_extract(rctx, *, whl_path, logger):
                 logger.debug(lambda: "Renaming: {} -> {}".format(src, dest))
                 repo_utils.rename(rctx, src, dest)
 
+        # Move RECORD to rewrite-record so gen_wheel_record can generate
+        # the platform-specific RECORD file at build time.
+        record_file = dist_info_dir.get_child("RECORD")
+        if record_file.exists:
+            rewrite_record_dir = rctx.path("rewrite-record/" + dist_info_dir.basename)
+            repo_utils.mkdir(rctx, rewrite_record_dir)
+            repo_utils.rename(rctx, record_file, rewrite_record_dir.get_child("RECORD"))
+
         # Ensure that there is no data dir left
         rctx.delete(data_dir)
 
+# Mapping of wheel .data categories to their extraction destination (relative to
+# repository root).
+_DATA_CATEGORIES = {
+    # category: repo_dest_dir
+    "data": "data",
+    "headers": "include",
+    # In theory there may be directory collisions in platlib/purelib, so it is
+    # best to merge the paths here. What is more, this code has to be reasonably
+    # efficient because some packages like to explicitly indicate if something
+    # is in `platlib` or `purelib` (e.g. libclang wheel).
+    "platlib": "site-packages",
+    "purelib": "site-packages",
+    "scripts": "bin",
+}
+
 def merge_trees(src, dest):
     """Merge src into the destination path.
 

python/private/pypi/whl_library_targets.bzl

@@ -19,6 +19,7 @@ load("//python:py_binary.bzl", "py_binary")
 load("//python:py_library.bzl", "py_library")
 load("//python/private:normalize_name.bzl", "normalize_name")
 load(":env_marker_setting.bzl", "env_marker_setting")
+load(":gen_wheel_record.bzl", "gen_wheel_record")
 load(
     ":labels.bzl",
     "DATA_LABEL",
@@ -159,6 +160,7 @@ def whl_library_srcs(
             py_library = py_library,
             venv_entry_point = venv_entry_point,
             venv_rewrite_shebang = venv_rewrite_shebang,
+            gen_wheel_record = gen_wheel_record,
             env_marker_setting = env_marker_setting,
             create_inits = _create_inits,
         )):
@@ -225,6 +227,16 @@ def whl_library_srcs(
         bins_for_data_label.append(rewrite_target_name)
         data.append(rewrite_target_name)
 
+    record_srcs = native.glob(["rewrite-record/*/RECORD"], allow_empty = True)
+    record_target_name = "record"
+    if record_srcs:
+        rules.gen_wheel_record(
+            name = record_target_name,
+            srcs = record_srcs,
+            tags = ["manual"],
+        )
+        data.append(record_target_name)
+
     if filegroups == None:
         filegroups = {
             EXTRACTED_WHEEL_FILES: dict(
@@ -248,6 +260,8 @@ def whl_library_srcs(
         srcs = native.glob(**glob_kwargs)
         if filegroup_name == DATA_LABEL:
             srcs = srcs + bins_for_data_label
+        if filegroup_name == DIST_INFO_LABEL and record_srcs:
+            srcs = srcs + [record_target_name]
         native.filegroup(
             name = filegroup_name,
             srcs = srcs,

tests/pypi/whl_extract/BUILD.bazel

@@ -0,0 +1,11 @@
+load("@rules_shell//shell:sh_test.bzl", "sh_test")
+load(":whl_extract_tests.bzl", "whl_extract_test_suite")
+
+whl_extract_test_suite(name = "whl_extract_tests")
+
+sh_test(
+    name = "wheel_record_rewriter_test",
+    srcs = ["wheel_record_rewriter_test.sh"],
+    args = ["$(location //python/private/pypi:wheel_record_rewriter)"],
+    data = ["//python/private/pypi:wheel_record_rewriter"],
+)