From ffc3176e05ad9fcb93bde1618b0923be1a7cc17f Mon Sep 17 00:00:00 2001
From: Richard Levasseur <rlevasseur@google.com>
Date: Fri, 15 Nov 2024 12:56:08 -0800
Subject: [PATCH] fix: keep imports if builtin py info missing

---
 CHANGELOG.md              |  3 ++-
 python/private/common.bzl | 25 ++++++++++++++++---------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c351487605..1acb12fb4d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -56,7 +56,8 @@ Unreleased changes template.
 
 {#v0-0-0-fixed}
 ### Fixed
-* Nothing yet.
+* (rules) Don't drop custom import paths if Bazel-builtin PyInfo is removed.
+  ([2414](https://github.com/bazelbuild/rules_python/issues/2414)).
 
 {#v0-0-0-added}
 ### Added
diff --git a/python/private/common.bzl b/python/private/common.bzl
index 837cd19097..97fabcebcb 100644
--- a/python/private/common.bzl
+++ b/python/private/common.bzl
@@ -263,15 +263,22 @@ def filter_to_py_srcs(srcs):
     return [f for f in srcs if f.extension == "py"]
 
 def collect_imports(ctx, semantics):
-    return depset(direct = semantics.get_imports(ctx), transitive = [
-        dep[PyInfo].imports
-        for dep in ctx.attr.deps
-        if PyInfo in dep
-    ] + [
-        dep[BuiltinPyInfo].imports
-        for dep in ctx.attr.deps
-        if BuiltinPyInfo in dep
-    ] if BuiltinPyInfo != None else [])
+    """Collect the direct and transitive `imports` strings.
+
+    Args:
+        ctx: {type}`ctx` the current target ctx
+        semantics: semantics object for fetching direct imports.
+
+    Returns:
+        {type}`depset[str]` of import paths
+    """
+    transitive = []
+    for dep in ctx.attr.deps:
+        if PyInfo in dep:
+            transitive.append(dep[PyInfo].imports)
+        if BuiltinPyInfo != None and BuiltinPyInfo in dep:
+            transitive.append(dep[BuiltinPyInfo].imports)
+    return depset(direct = semantics.get_imports(ctx), transitive = transitive)
 
 def collect_runfiles(ctx, files = depset()):
     """Collects the necessary files from the rule's context.