Forward new provider for python in filter_layer

lang/image.bzl

@@ -297,17 +297,30 @@ def _filter_layer_rule_impl(ctx):
         if str(dep.target.label).startswith(ctx.attr.filter) and str(dep.target.label) != str(ctx.attr.dep.label):
             runfiles = runfiles.merge(dep.target.default_runfiles)
             filtered_depsets.append(dep.target_deps)
-    return struct(
-        providers = [
-            FilterLayerInfo(
-                runfiles = runfiles,
-                filtered_depset = depset(transitive = filtered_depsets),
-            ),
-        ],
-        # Also forward builtin providers so that the filter_layer() can be used as a normal
-        # dependency to native targets (e.g. py_library(deps = [<filter_layer>])).
-        py = ctx.attr.dep.py if hasattr(ctx.attr.dep, "py") else None,
-    )
+
+    # Forward correct provider, depending on availability, so that the filter_layer() can be
+    # used as a normal dependency to native targets (e.g. py_library(deps = [<filter_layer>])).
+    if hasattr(ctx.attr.dep, "py"):
+        # Forward legacy builtin provider and PyInfo provider
+        return struct(
+            providers = [
+                FilterLayerInfo(
+                    runfiles = runfiles,
+                    filtered_depset = depset(transitive = filtered_depsets),
+                ),
+            ] + ([ctx.attr.dep[PyInfo]] if PyInfo in ctx.attr.dep else []),
+            py = ctx.attr.dep.py if hasattr(ctx.attr.dep, "py") else None,
+        )
+    else:
+        # Forward the PyInfo provider only
+        return struct(
+            providers = [
+                FilterLayerInfo(
+                    runfiles = runfiles,
+                    filtered_depset = depset(transitive = filtered_depsets),
+                ),
+            ] + ([ctx.attr.dep[PyInfo]] if PyInfo in ctx.attr.dep else []),
+        )
 
 # A rule that allows selecting a subset of transitive dependencies, and using
 # them as a layer in an image.