Adding variable support for distribution in py_wheel

ns-tkonduri · ns-tkonduri · commit 385c10303e47 · 2023-06-09T16:45:12.000-07:00
diff --git a/docs/packaging.md b/docs/packaging.md
diff --git a/examples/wheel/BUILD.bazel b/examples/wheel/BUILD.bazel
@@ -94,7 +94,7 @@ build_test(
 py_wheel(
     name = "minimal_with_py_library_with_stamp",
     # Package data. We're building "example_minimal_library-0.0.1-py3-none-any.whl"
-    distribution = "example_minimal_library",
+    distribution = "example_minimal_library{BUILD_USER}",
     python_tag = "py3",
     stamp = 1,
     version = "0.1.{BUILD_TIMESTAMP}",
diff --git a/examples/wheel/wheel_test.py b/examples/wheel/wheel_test.py
@@ -374,30 +374,36 @@ def test_rule_creates_directory_and_is_included_in_wheel(self):
                 ],
             )
 
-    def test_rule_sets_stamped_version_in_wheel_metadata(self):
+    def test_rule_expands_workspace_status_keys_in_wheel_metadata(self):
         filename = os.path.join(
             os.environ["TEST_SRCDIR"],
             "rules_python",
             "examples",
             "wheel",
-            "example_minimal_library-0.1._BUILD_TIMESTAMP_-py3-none-any.whl",
+            "example_minimal_library_BUILD_USER_-0.1._BUILD_TIMESTAMP_-py3-none-any.whl",
         )
 
         with zipfile.ZipFile(filename) as zf:
             metadata_file = None
             for f in zf.namelist():
                 self.assertNotIn("_BUILD_TIMESTAMP_", f)
+                self.assertNotIn("_BUILD_USER_", f)
                 if os.path.basename(f) == "METADATA":
                     metadata_file = f
             self.assertIsNotNone(metadata_file)
 
             version = None
+            name = None
             with zf.open(metadata_file) as fp:
                 for line in fp:
-                    if line.startswith(b'Version:'):
+                    if line.startswith(b"Version:"):
                         version = line.decode().split()[-1]
+                    if line.startswith(b"Name:"):
+                        name = line.decode().split()[-1]
             self.assertIsNotNone(version)
+            self.assertIsNotNone(name)
             self.assertNotIn("{BUILD_TIMESTAMP}", version)
+            self.assertNotIn("{BUILD_USER}", name)
 
 
 if __name__ == "__main__":
diff --git a/python/private/py_wheel.bzl b/python/private/py_wheel.bzl
@@ -40,6 +40,12 @@ Name of the distribution.
 
 This should match the project name onm PyPI. It's also the name that is used to
 refer to the package in other packages' dependencies.
+
+Workspace status keys are expanded using `{NAME}` format, for example:
+ - `distribution = "package.{CLASSIFIER}"`
+ - `distribution = "{DISTRIBUTION}"`
+
+For the available keys, see https://bazel.build/docs/user-manual#workspace-status
 """,
     ),
     "platform": attr.string(
diff --git a/tools/wheelmaker.py b/tools/wheelmaker.py
@@ -167,12 +167,12 @@ def add_wheelfile(self):
             wheel_contents += "Tag: %s\n" % tag
         self.add_string(self.distinfo_path("WHEEL"), wheel_contents)
 
-    def add_metadata(self, metadata, description, version):
+    def add_metadata(self, metadata, name, description, version):
         """Write METADATA file to the distribution."""
         # https://www.python.org/dev/peps/pep-0566/
         # https://packaging.python.org/specifications/core-metadata/
-        metadata += "Version: " + version
-        metadata += "\n\n"
+        metadata = re.sub("^Name: .*$", "Name: %s" % name, metadata, flags=re.MULTILINE)
+        metadata += "Version: %s\n\n" % version
         # setuptools seems to insert UNKNOWN as description when none is
         # provided.
         metadata += description if description else "UNKNOWN"
@@ -207,18 +207,18 @@ def get_files_to_package(input_files):
     return files
 
 
-def resolve_version_stamp(
-    version: str, volatile_status_stamp: Path, stable_status_stamp: Path
+def resolve_argument_stamp(
+    argument: str, volatile_status_stamp: Path, stable_status_stamp: Path
 ) -> str:
-    """Resolve workspace status stamps format strings found in the version string
+    """Resolve workspace status stamps format strings found in the argument string
 
     Args:
-        version (str): The raw version represenation for the wheel (may include stamp variables)
+        argument (str): The raw argument represenation for the wheel (may include stamp variables)
         volatile_status_stamp (Path): The path to a volatile workspace status file
         stable_status_stamp (Path): The path to a stable workspace status file
 
     Returns:
-        str: A resolved version string
+        str: A resolved argument string
     """
     lines = (
         volatile_status_stamp.read_text().splitlines()
@@ -229,9 +229,9 @@ def resolve_version_stamp(
             continue
         key, value = line.split(" ", maxsplit=1)
         stamp = "{" + key + "}"
-        version = version.replace(stamp, value)
+        argument = argument.replace(stamp, value)
 
-    return version
+    return argument
 
 
 def parse_args() -> argparse.Namespace:
@@ -357,7 +357,16 @@ def main() -> None:
     strip_prefixes = [p for p in arguments.strip_path_prefix]
 
     if arguments.volatile_status_file and arguments.stable_status_file:
-        version = resolve_version_stamp(
+        name = resolve_argument_stamp(
+            arguments.name,
+            arguments.volatile_status_file,
+            arguments.stable_status_file,
+        )
+    else:
+        name = arguments.name
+
+    if arguments.volatile_status_file and arguments.stable_status_file:
+        version = resolve_argument_stamp(
             arguments.version,
             arguments.volatile_status_file,
             arguments.stable_status_file,
@@ -366,7 +375,7 @@ def main() -> None:
         version = arguments.version
 
     with WheelMaker(
-        name=arguments.name,
+        name=name,
         version=version,
         build_tag=arguments.build_tag,
         python_tag=arguments.python_tag,
@@ -398,7 +407,9 @@ def main() -> None:
             with open(arguments.metadata_file, "rt", encoding="utf-8") as metadata_file:
                 metadata = metadata_file.read()
 
-        maker.add_metadata(metadata=metadata, description=description, version=version)
+        maker.add_metadata(
+            metadata=metadata, name=name, description=description, version=version
+        )
 
         if arguments.entry_points_file:
             maker.add_file(
