fix: ensure transitive wheels are passed from py_library

internal_python_deps.bzl

@@ -6,11 +6,30 @@ statement from these, that's a bug in our distribution.
 These happen after the regular internal dependencies loads as we need to reference the resolved interpreter
 """
 
-load("@rules_python//python:pip.bzl", "pip_parse")
+load("@rules_python//python:pip.bzl", "package_annotation", "pip_parse")
+
+PY_WHEEL_RULE_CONTENT = """\
+load("@aspect_rules_py//py:defs.bzl", "py_wheel")
+py_wheel(
+    name = "wheel",
+    src = ":whl",
+)
+"""
 
 def rules_py_internal_pypi_deps(interpreter):
+    # Here we can see an example of annotations being applied to an arbitrary
+    # package. For details on `package_annotation` and it's uses, see the
+    # docs at @rules_python//docs:pip.md`.
+
+    PACKAGES = ["django", "colorama"]
+    ANNOTATIONS = {
+        pkg: package_annotation(additive_build_content = PY_WHEEL_RULE_CONTENT)
+        for pkg in PACKAGES
+    }
+
     pip_parse(
         name = "pypi",
+        annotations = ANNOTATIONS,
         python_interpreter_target = interpreter,
         requirements_lock = "//py/tests/external-deps:requirements.txt",
     )

py/private/BUILD.bazel

@@ -32,21 +32,23 @@ bzl_library(
     visibility = ["//:__subpackages__"],
     deps = [
         ":providers",
+        ":py_wheel",
         "@bazel_skylib//lib:paths",
+        "@bazel_skylib//lib:types",
     ],
 )
 
 bzl_library(
-    name = "providers",
-    srcs = ["providers.bzl"],
+    name = "py_wheel",
+    srcs = ["py_wheel.bzl"],
     visibility = ["//py:__subpackages__"],
+    deps = [":providers"],
 )
 
 bzl_library(
-    name = "py_wheel",
-    srcs = ["py_wheel.bzl"],
+    name = "providers",
+    srcs = ["providers.bzl"],
     visibility = ["//py:__subpackages__"],
-    deps = [":providers"],
 )
 
 bzl_library(

py/private/py_library.bzl

@@ -2,6 +2,7 @@
 
 load("@bazel_skylib//lib:paths.bzl", "paths")
 load("//py/private:providers.bzl", "PyWheelInfo")
+load("//py/private:py_wheel.bzl", py_wheel = "py_wheel_lib")
 
 def _make_srcs_depset(ctx):
     return depset(
@@ -56,6 +57,7 @@ def _py_library_impl(ctx):
     transitive_srcs = _make_srcs_depset(ctx)
     imports = _make_imports_depset(ctx)
     runfiles = _make_merged_runfiles(ctx)
+    py_wheel_info = py_wheel.make_py_wheel_info(ctx, ctx.attr.deps)
 
     return [
         DefaultInfo(
@@ -69,6 +71,7 @@ def _py_library_impl(ctx):
             has_py3_only_sources = True,
             uses_shared_libraries = False,
         ),
+        py_wheel_info,
     ]
 
 _attrs = dict({

py/private/py_wheel.bzl

@@ -1,5 +1,6 @@
 "Represent a python wheel file"
 
+load("@bazel_skylib//lib:types.bzl", "types")
 load("//py/private:providers.bzl", "PyWheelInfo")
 
 _ATTRS = {
@@ -8,18 +9,34 @@ _ATTRS = {
     ),
 }
 
-def _make_py_wheel_info_from_filegroup(wheel_filegroup):
+def _make_py_wheel_info(ctx, wheel_filegroups):
+    if not types.is_list(wheel_filegroups):
+        filegroups = [wheel_filegroups]
+    else:
+        filegroups = wheel_filegroups
+
     files_depsets = []
-    files_depsets.append(wheel_filegroup[DefaultInfo].files)
-    files_depsets.append(wheel_filegroup[DefaultInfo].default_runfiles.files)
+    runfiles = []
+    for filegroup in filegroups:
+        if DefaultInfo in filegroup:
+            files_depsets.append(filegroup[DefaultInfo].files)
+            files_depsets.append(filegroup[DefaultInfo].default_runfiles.files)
+            runfiles.append(filegroup[DefaultInfo].default_runfiles)
+
+        if PyWheelInfo in filegroup:
+            files_depsets.append(filegroup[PyWheelInfo].files)
+            runfiles.append(filegroup[PyWheelInfo].default_runfiles)
+
+    py_info_runfiles = ctx.runfiles()
+    py_info_runfiles = py_info_runfiles.merge_all(runfiles)
 
     return PyWheelInfo(
         files = depset(transitive = files_depsets),
-        default_runfiles = wheel_filegroup[DefaultInfo].default_runfiles,
+        default_runfiles = py_info_runfiles,
     )
 
 def _py_wheel_impl(ctx):
-    py_wheel_info = _make_py_wheel_info_from_filegroup(ctx.attr.src)
+    py_wheel_info = _make_py_wheel_info(ctx, ctx.attr.src)
     return [
         py_wheel_info,
     ]
@@ -28,4 +45,5 @@ py_wheel_lib = struct(
     implementation = _py_wheel_impl,
     attrs = _ATTRS,
     provides = [PyWheelInfo],
+    make_py_wheel_info = _make_py_wheel_info,
 )

py/tests/external-deps/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@rules_python//python/pip_install:requirements.bzl", "compile_pip_requirements")
-load("//py:defs.bzl", "py_binary", "py_library", "py_wheel")
+load("//py:defs.bzl", "py_binary", "py_library")
 load("@python_toolchain//:defs.bzl", "host_platform")
 load("@aspect_bazel_lib//lib:write_source_files.bzl", "write_source_files")
 
@@ -9,22 +9,20 @@ compile_pip_requirements(
     requirements_txt = "requirements.txt",
 )
 
-py_wheel(
-    name = "django",
-    src = "@pypi_django//:whl",
-)
-
 py_library(
     name = "lib",
     srcs = ["lib.py"],
+    deps = [
+        "@pypi_colorama//:wheel",
+    ],
 )
 
 py_binary(
     name = "main",
     srcs = ["__main__.py"],
     deps = [
-        ":django",
         ":lib",
+        "@pypi_django//:wheel",
     ],
 )
 
@@ -35,6 +33,7 @@ genrule(
         "$(execpath main)",
         """sed "s#$$(pwd)#(pwd)#" """,
         "sed 's#^.*execroot/aspect_rules_py/external/python_toolchain_%s#(py_toolchain)#'" % host_platform,
+        "sed 's#(main, .*)#(main, REDACTED)#'",
     ]) + "> $(execpath out)",
     tools = ["main"],
 )

py/tests/external-deps/__main__.py

@@ -2,6 +2,7 @@
 import site
 import sys
 import django
+import inspect
 
 print(f'Python: {sys.executable}')
 print(f'version: {sys.version}')
@@ -19,4 +20,5 @@
 print(f'Django version: {django.__version__}')
 
 from lib import greet
-print(greet("Matt"))
+print(f'\nFrom lib with wheel dependency: {greet("Matt")}')
+print(f'lib filepath: {inspect.getsourcefile(greet)}')

py/tests/external-deps/expected

@@ -1,5 +1,5 @@
 Python: (pwd)/bazel-out/host/bin/py/tests/external-deps/main.runfiles/.main.venv/bin/python3
-version: 3.9.10 (main, Feb 27 2022, 23:39:20) 
+version: 3.9.10 (main, REDACTED) 
 [Clang 13.0.1 ]
 version info: sys.version_info(major=3, minor=9, micro=10, releaselevel='final', serial=0)
 cwd: (pwd)
@@ -16,4 +16,6 @@ Entrypoint Path: (pwd)/bazel-out/host/bin/py/tests/external-deps/main.runfiles/a
 
 Django location: (pwd)/bazel-out/host/bin/py/tests/external-deps/main.runfiles/.main.venv/lib/python3.9/site-packages/django/__init__.py
 Django version: 4.0.2
-Hello Matt
+
+From lib with wheel dependency: Hello Matt
+lib filepath: (pwd)/bazel-out/host/bin/py/tests/external-deps/main.runfiles/aspect_rules_py/py/tests/external-deps/lib.py

py/tests/external-deps/lib.py

@@ -1,2 +1,4 @@
+from colorama import Fore, Style
+
 def greet(name):
-    return f"Hello {name}"
+    return f"{Fore.GREEN}Hello {name}{Style.RESET_ALL}"

py/tests/external-deps/requirements.in

@@ -1 +1,2 @@
-django==4.0.2
+django==4.0.2
+colorama==0.4.4

py/tests/external-deps/requirements.txt

@@ -8,6 +8,10 @@ asgiref==3.5.0 \
     --hash=sha256:2f8abc20f7248433085eda803936d98992f1343ddb022065779f37c5da0181d0 \
     --hash=sha256:88d59c13d634dcffe0510be048210188edd79aeccb6a6c9028cdad6f31d730a9
     # via django
+colorama==0.4.4 \
+    --hash=sha256:5941b2b48a20143d2267e95b1c2a7603ce057ee39fd88e7329b0c292aa16869b \
+    --hash=sha256:9f47eda37229f68eee03b24b9748937c7dc3868f906e8ba69fbcbdd3bc5dc3e2
+    # via -r py/tests/external-deps/requirements.in
 django==4.0.2 \
     --hash=sha256:110fb58fb12eca59e072ad59fc42d771cd642dd7a2f2416582aa9da7a8ef954a \
     --hash=sha256:996495c58bff749232426c88726d8cd38d24c94d7c1d80835aafffa9bc52985a