From cdd69dbe86f0ee53626816b72faa65db6c8c93d9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 21 Oct 2024 09:45:27 -0700
Subject: [PATCH] Add macros for unpacking wheels and using the wheels content
 in py_library.

PiperOrigin-RevId: 688169758
---
 opensource_only.files                   |  1 +
 third_party/py/python_wheel_library.bzl | 39 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+)
 create mode 100644 third_party/py/python_wheel_library.bzl

diff --git a/opensource_only.files b/opensource_only.files
index b96a4bc89..265b15c67 100644
--- a/opensource_only.files
+++ b/opensource_only.files
@@ -108,6 +108,7 @@ third_party/py/python_init_repositories.bzl:
 third_party/py/python_init_rules.bzl:
 third_party/py/python_init_toolchains.bzl:
 third_party/py/python_repo.bzl:
+third_party/py/python_wheel_library.bzl:
 third_party/pybind11.BUILD:
 third_party/pybind11_bazel/BUILD:
 third_party/python_runtime/BUILD:
diff --git a/third_party/py/python_wheel_library.bzl b/third_party/py/python_wheel_library.bzl
new file mode 100644
index 000000000..c949d8800
--- /dev/null
+++ b/third_party/py/python_wheel_library.bzl
@@ -0,0 +1,39 @@
+""" Macros to unpack a wheel and use its content as a py_library. """
+
+def _unpacked_wheel_impl(ctx):
+    output_dir = ctx.actions.declare_directory(ctx.label.name)
+    ctx.actions.run(
+        inputs = [ctx.file.wheel],
+        outputs = [output_dir],
+        executable = ctx.executable.zipper,
+        arguments = ["x", ctx.file.wheel.path, "-d", output_dir.path],
+    )
+    return [
+        DefaultInfo(files = depset([output_dir])),
+    ]
+
+_unpacked_wheel = rule(
+    implementation = _unpacked_wheel_impl,
+    attrs = {
+        "wheel": attr.label(mandatory = True, allow_single_file = True),
+        "zipper": attr.label(
+            default = Label("@bazel_tools//tools/zip:zipper"),
+            cfg = "exec",
+            executable = True,
+        ),
+    },
+)
+
+def wheel_library(name, wheel, deps):
+    unpacked_wheel_name = name + "_unpacked_wheel"
+    _unpacked_wheel(
+        name = unpacked_wheel_name,
+        wheel = wheel,
+    )
+    native.py_library(
+        name = name,
+        data = [":" + unpacked_wheel_name],
+        imports = [unpacked_wheel_name],
+        deps = deps,
+        visibility = ["//visibility:public"],
+    )