From a128bdbb2c28ed3e31f6988271d7124d84ff0b2d Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Thu, 24 Oct 2024 14:11:31 -0700
Subject: [PATCH] Add TF wheel compliance verification in Bazel build rule.

The compliance check is disabled by default, but can be turned on by `--//tensorflow/tools/pip_package:wheel_compliance=true` option.

PiperOrigin-RevId: 689511666
---
 .../tsl/third_party/py/python_wheel_library.bzl     | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/third_party/tsl/third_party/py/python_wheel_library.bzl b/third_party/tsl/third_party/py/python_wheel_library.bzl
index 01ea7554e2323a..c2d2965addbf66 100644
--- a/third_party/tsl/third_party/py/python_wheel_library.bzl
+++ b/third_party/tsl/third_party/py/python_wheel_library.bzl
@@ -9,6 +9,11 @@ def _unpacked_wheel_impl(ctx):
             if linker_input.libraries and linker_input.libraries[0].dynamic_library:
                 lib = linker_input.libraries[0].dynamic_library
                 libs.append(lib)
+    wheel = [
+        w
+        for w in ctx.files.wheel_rule_outputs
+        if w.basename.endswith(".whl")
+    ][0]
     script = """
     {zipper} x {wheel} -d {output}
     for lib in {libs}; do
@@ -16,12 +21,12 @@ def _unpacked_wheel_impl(ctx):
     done
     """.format(
         zipper = ctx.executable.zipper.path,
-        wheel = ctx.file.wheel.path,
+        wheel = wheel.path,
         output = output_dir.path,
         libs = " ".join(["'%s'" % lib.path for lib in libs]),
     )
     ctx.actions.run_shell(
-        inputs = [ctx.file.wheel] + libs,
+        inputs = ctx.files.wheel_rule_outputs + libs,
         command = script,
         outputs = [output_dir],
         tools = [ctx.executable.zipper],
@@ -34,7 +39,7 @@ def _unpacked_wheel_impl(ctx):
 _unpacked_wheel = rule(
     implementation = _unpacked_wheel_impl,
     attrs = {
-        "wheel": attr.label(mandatory = True, allow_single_file = True),
+        "wheel_rule_outputs": attr.label(mandatory = True, allow_files = True),
         "zipper": attr.label(
             default = Label("@bazel_tools//tools/zip:zipper"),
             cfg = "exec",
@@ -48,7 +53,7 @@ def wheel_library(name, wheel, deps = [], wheel_deps = []):
     unpacked_wheel_name = name + "_unpacked_wheel"
     _unpacked_wheel(
         name = unpacked_wheel_name,
-        wheel = wheel,
+        wheel_rule_outputs = wheel,
         deps = wheel_deps,
     )
     native.py_library(