From 946932144de329607a3d38f8204e76cf0e3ed57e Mon Sep 17 00:00:00 2001
From: "Holmes, Brad" <holmesbr@ncbi.nlm.nih.gov>
Date: Tue, 12 Mar 2024 16:21:05 -0400
Subject: [PATCH] feat: Add support for python-wheel data directory

This feature addes  attribute to the py_wheel rule and
makes a related change to the wheelmaker tool.

Fixes https://github.com/bazelbuild/rules_python/issues/1777
---
 python/private/repack_whl.py      |  3 ++-
 tests/py_wheel/py_wheel_tests.bzl | 25 +++++++++++++++++++++++++
 tools/wheelmaker.py               |  2 +-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/python/private/repack_whl.py b/python/private/repack_whl.py
index be113ef791..ea9c01f76f 100644
--- a/python/private/repack_whl.py
+++ b/python/private/repack_whl.py
@@ -150,8 +150,9 @@ def main(sys_argv):
         logging.debug(f"Found dist-info dir: {distinfo_dir}")
         record_path = distinfo_dir / "RECORD"
         record_contents = record_path.read_text() if record_path.exists() else ""
+        distribution_prefix = distinfo_dir.with_suffix("").name
 
-        with _WhlFile(args.output, mode="w", distinfo_dir=distinfo_dir) as out:
+        with _WhlFile(args.output, mode="w", distribution_prefix=distribution_prefix) as out:
             for p in _files_to_pack(patched_wheel_dir, record_contents):
                 rel_path = p.relative_to(patched_wheel_dir)
                 out.add_file(str(rel_path), p)
diff --git a/tests/py_wheel/py_wheel_tests.bzl b/tests/py_wheel/py_wheel_tests.bzl
index 3c03a1b8e4..6ece127d59 100644
--- a/tests/py_wheel/py_wheel_tests.bzl
+++ b/tests/py_wheel/py_wheel_tests.bzl
@@ -46,6 +46,31 @@ def _test_metadata_impl(env, target):
 
 _tests.append(_test_metadata)
 
+def _test_data(name):
+    rt_util.helper_target(
+        py_wheel,
+        name = name + "_data",
+        distribution = "mydist_" + name,
+        version = "0.0.0",
+        data_files = {
+            "source_name": "wheel_name",
+        },
+    )
+    analysis_test(
+        name = name,
+        impl = _test_data_impl,
+        target = name + "_data",
+    )
+
+def _test_data_impl(env, target):
+    action = env.expect.that_target(target).action_named(
+        "PyWheel",
+    )
+    action.contains_at_least_args(["--data_files", "wheel_name;tests/py_wheel/source_name"])
+    action.contains_at_least_inputs(["tests/py_wheel/source_name"])
+
+_tests.append(_test_data)
+
 def _test_content_type_from_attr(name):
     rt_util.helper_target(
         py_wheel,
diff --git a/tools/wheelmaker.py b/tools/wheelmaker.py
index 93f9cb5514..aef897b940 100644
--- a/tools/wheelmaker.py
+++ b/tools/wheelmaker.py
@@ -102,7 +102,7 @@ def __init__(
         filename,
         *,
         mode,
-        distribution_prefix: str | Path,
+        distribution_prefix: str,
         strip_path_prefixes=None,
         compression=zipfile.ZIP_DEFLATED,
         **kwargs,