1 files changed, 75 insertions, 0 deletions
diff --git a/python/private/py_package.bzl b/python/private/py_package.bzl
new file mode 100644
index 0000000..08f4b0b
--- /dev/null
+++ b/python/private/py_package.bzl
@@ -0,0 +1,75 @@
+# Copyright 2023 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"Implementation of py_package rule"
+
+def _path_inside_wheel(input_file):
+    # input_file.short_path is sometimes relative ("../${repository_root}/foobar")
+    # which is not a valid path within a zip file. Fix that.
+    short_path = input_file.short_path
+    if short_path.startswith("..") and len(short_path) >= 3:
+        # Path separator. '/' on linux.
+        separator = short_path[2]
+
+        # Consume '../' part.
+        short_path = short_path[3:]
+
+        # Find position of next '/' and consume everything up to that character.
+        pos = short_path.find(separator)
+        short_path = short_path[pos + 1:]
+    return short_path
+
+def _py_package_impl(ctx):
+    inputs = depset(
+        transitive = [dep[DefaultInfo].data_runfiles.files for dep in ctx.attr.deps] +
+                     [dep[DefaultInfo].default_runfiles.files for dep in ctx.attr.deps],
+    )
+
+    # TODO: '/' is wrong on windows, but the path separator is not available in starlark.
+    # Fix this once ctx.configuration has directory separator information.
+    packages = [p.replace(".", "/") for p in ctx.attr.packages]
+    if not packages:
+        filtered_inputs = inputs
+    else:
+        filtered_files = []
+
+        # TODO: flattening depset to list gives poor performance,
+        for input_file in inputs.to_list():
+            wheel_path = _path_inside_wheel(input_file)
+            for package in packages:
+                if wheel_path.startswith(package):
+                    filtered_files.append(input_file)
+        filtered_inputs = depset(direct = filtered_files)
+
+    return [DefaultInfo(
+        files = filtered_inputs,
+    )]
+
+py_package_lib = struct(
+    implementation = _py_package_impl,
+    attrs = {
+        "deps": attr.label_list(
+            doc = "",
+        ),
+        "packages": attr.string_list(
+            mandatory = False,
+            allow_empty = True,
+            doc = """\
+List of Python packages to include in the distribution.
+Sub-packages are automatically included.
+""",
+        ),
+    },
+    path_inside_wheel = _path_inside_wheel,
+)