From 6d8c97474af4cb67a7e68810172f471ae2dbeae5 Mon Sep 17 00:00:00 2001
From: Matt Mackay <mattem@gmail.com>
Date: Wed, 25 May 2022 18:44:04 -0400
Subject: [PATCH] fix: don't treat all transitive py_library rules as external

---
 py/private/py_wheel.bzl            | 12 +++---
 py/tests/internal-deps/BUILD       | 63 ++++++++++++++++++++++++++++++
 py/tests/internal-deps/__init__.py |  0
 py/tests/internal-deps/__main__.py |  7 ++++
 py/tests/internal-deps/add.py      |  2 +
 py/tests/internal-deps/expected    |  3 ++
 py/tests/internal-deps/pi.py       |  4 ++
 py/tests/internal-deps/sub.py      |  2 +
 8 files changed, 88 insertions(+), 5 deletions(-)
 create mode 100644 py/tests/internal-deps/BUILD
 create mode 100644 py/tests/internal-deps/__init__.py
 create mode 100644 py/tests/internal-deps/__main__.py
 create mode 100644 py/tests/internal-deps/add.py
 create mode 100644 py/tests/internal-deps/expected
 create mode 100644 py/tests/internal-deps/pi.py
 create mode 100644 py/tests/internal-deps/sub.py

diff --git a/py/private/py_wheel.bzl b/py/private/py_wheel.bzl
index 2d208f16..604a0939 100644
--- a/py/private/py_wheel.bzl
+++ b/py/private/py_wheel.bzl
@@ -18,14 +18,16 @@ def _make_py_wheel_info(ctx, wheel_filegroups):
     files_depsets = []
     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)
-
+        # The ordering is important here as we want to ensure we use the PyWheelInfo from transitive
+        # py_library dependencies, and only fall back to DefaultInfo when translating from the wheel
+        # filegroup to py_wheel_library
         if PyWheelInfo in filegroup:
             files_depsets.append(filegroup[PyWheelInfo].files)
             runfiles.append(filegroup[PyWheelInfo].default_runfiles)
+        elif DefaultInfo in filegroup:
+            files_depsets.append(filegroup[DefaultInfo].files)
+            files_depsets.append(filegroup[DefaultInfo].default_runfiles.files)
+            runfiles.append(filegroup[DefaultInfo].default_runfiles)
 
     py_info_runfiles = ctx.runfiles()
     py_info_runfiles = py_info_runfiles.merge_all(runfiles)
diff --git a/py/tests/internal-deps/BUILD b/py/tests/internal-deps/BUILD
new file mode 100644
index 00000000..f043600f
--- /dev/null
+++ b/py/tests/internal-deps/BUILD
@@ -0,0 +1,63 @@
+load("//py:defs.bzl", "py_binary", "py_test", rules_py_py_library = "py_library")
+load("@aspect_bazel_lib//lib:write_source_files.bzl", "write_source_files")
+
+rules_py_py_library(
+    name = "init",
+    srcs = [
+        "__init__.py",
+    ],
+)
+
+rules_py_py_library(
+    name = "add",
+    srcs = [
+        "add.py",
+    ],
+    deps = [
+        ":init",
+    ],
+)
+
+rules_py_py_library(
+    name = "sub",
+    srcs = [
+        "sub.py",
+    ],
+    deps = [
+        ":init",
+    ],
+)
+
+py_library(
+    name = "pi",
+    srcs = [
+        "pi.py",
+    ],
+    deps = [
+        ":init",
+    ],
+)
+
+py_binary(
+    name = "main",
+    srcs = [
+        "__main__.py",
+    ],
+    deps = [
+        ":add",
+        ":pi",
+        ":sub",
+    ],
+)
+
+py_test(
+    name = "assert",
+    srcs = [
+        "__main__.py",
+    ],
+    deps = [
+        ":add",
+        ":pi",
+        ":sub",
+    ],
+)
diff --git a/py/tests/internal-deps/__init__.py b/py/tests/internal-deps/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/py/tests/internal-deps/__main__.py b/py/tests/internal-deps/__main__.py
new file mode 100644
index 00000000..ab77949d
--- /dev/null
+++ b/py/tests/internal-deps/__main__.py
@@ -0,0 +1,7 @@
+from add import add
+from sub import sub
+from pi import format_pi
+
+assert add(3, .14) == 3.14
+assert sub(4, .86) == 3.14
+assert format_pi('.2f') == "3.14"
diff --git a/py/tests/internal-deps/add.py b/py/tests/internal-deps/add.py
new file mode 100644
index 00000000..4693ad3c
--- /dev/null
+++ b/py/tests/internal-deps/add.py
@@ -0,0 +1,2 @@
+def add(a, b):
+    return a + b
diff --git a/py/tests/internal-deps/expected b/py/tests/internal-deps/expected
new file mode 100644
index 00000000..e1333864
--- /dev/null
+++ b/py/tests/internal-deps/expected
@@ -0,0 +1,3 @@
+3.14
+3.14
+3.14
diff --git a/py/tests/internal-deps/pi.py b/py/tests/internal-deps/pi.py
new file mode 100644
index 00000000..084e4e24
--- /dev/null
+++ b/py/tests/internal-deps/pi.py
@@ -0,0 +1,4 @@
+from math import pi
+
+def format_pi(f):
+    return format(pi, f)
\ No newline at end of file
diff --git a/py/tests/internal-deps/sub.py b/py/tests/internal-deps/sub.py
new file mode 100644
index 00000000..312b477d
--- /dev/null
+++ b/py/tests/internal-deps/sub.py
@@ -0,0 +1,2 @@
+def sub(a, b):
+    return a - b