feat(py): propagate type stubs and carry them into venvs

docs/api/py.md

@@ -170,7 +170,7 @@ py_library(<a href="#py_library-name">name</a>, <a href="#py_library-deps">deps<
 | :------------- | :------------- | :------------- | :------------- | :------------- |
 | <a id="py_library-name"></a>name |  A unique name for this target.   | <a href="https://bazel.build/concepts/labels#target-names">Name</a> | required |  |
 | <a id="py_library-deps"></a>deps |  Targets that produce Python code, commonly `py_library` rules.   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional |  `[]`  |
-| <a id="py_library-srcs"></a>srcs |  Python source files.   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional |  `[]`  |
+| <a id="py_library-srcs"></a>srcs |  Python source files.<br><br>`.pyi` type stubs listed here are carried as `PyInfo.transitive_pyi_files` rather than as runtime sources; both reach the runfiles of venvs and launchers that include this library.   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional |  `[]`  |
 | <a id="py_library-data"></a>data |  Runtime dependencies of the program.<br><br>The transitive closure of the `data` dependencies will be available in the `.runfiles` folder for this binary/test. The program may optionally use the Runfiles lookup library to locate the data files, see https://pypi.org/project/bazel-runfiles/. Data is analyzed in the inherited caller configuration. Put artifacts that must match the terminal's Python environment in `deps`.   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional |  `[]`  |
 | <a id="py_library-imports"></a>imports |  List of import directories to be added to the PYTHONPATH.   | List of strings | optional |  `[]`  |
 | <a id="py_library-resolutions"></a>resolutions |  Satisfy a virtual_dep with a mapping from external package name to the label of an installed package that provides it. See virtual_deps.   | Dictionary: String -> Label | optional |  `{}`  |
@@ -288,7 +288,7 @@ from the already extracted whl file.
 <pre>
 load("@aspect_rules_py//py:defs.bzl", "PyInfo")
 
-PyInfo(<a href="#PyInfo-transitive_sources">transitive_sources</a>, <a href="#PyInfo-imports">imports</a>, <a href="#PyInfo-virtual_dependencies">virtual_dependencies</a>, <a href="#PyInfo-virtual_resolutions">virtual_resolutions</a>)
+PyInfo(<a href="#PyInfo-transitive_sources">transitive_sources</a>, <a href="#PyInfo-transitive_pyi_files">transitive_pyi_files</a>, <a href="#PyInfo-imports">imports</a>, <a href="#PyInfo-virtual_dependencies">virtual_dependencies</a>, <a href="#PyInfo-virtual_resolutions">virtual_resolutions</a>)
 </pre>
 
 Python source, import-path, and virtual-dependency information for a target's dependency closure.
@@ -297,7 +297,8 @@ Python source, import-path, and virtual-dependency information for a target's de
 
 | Name  | Description |
 | :------------- | :------------- |
-| <a id="PyInfo-transitive_sources"></a>transitive_sources |  depset[File] — postorder depset of first-party `.py` sources in the transitive closure.    |
+| <a id="PyInfo-transitive_sources"></a>transitive_sources |  depset[File] — postorder depset of first-party runtime sources in the transitive closure; `.pyi` stubs are excluded.    |
+| <a id="PyInfo-transitive_pyi_files"></a>transitive_pyi_files |  depset[File] — postorder depset of `.pyi` type stubs in the transitive closure: stubs listed in `srcs` plus those carried by deps of either ruleset.    |
 | <a id="PyInfo-imports"></a>imports |  depset[str] — import roots to place on `sys.path` (rlocation-root-relative).    |
 | <a id="PyInfo-virtual_dependencies"></a>virtual_dependencies |  depset[str] — names of required virtual dependencies, independent of their resolution status.    |
 | <a id="PyInfo-virtual_resolutions"></a>virtual_resolutions |  depset[struct(virtual, target)] — virtual-dependency-name to concrete-target resolutions.    |

e2e/rules-python-interop/BUILD.bazel

@@ -8,6 +8,7 @@ load("@rules_python//python/zipapp:py_zipapp_binary.bzl", "py_zipapp_binary")
 load("@rules_shell//shell:sh_test.bzl", "sh_test")
 load(":exec_tools_facts.bzl", "exec_tools_facts", "with_python_version")
 load(":report_toolchain_version.bzl", "report_exec_version", "report_version")
+load(":type_stubs_test.bzl", "foreign_stubs_test")
 
 # Both rulesets provision interpreters here, split by version so neither
 # shadows the other (see MODULE.bazel): 3.11 is rules_python's, everything else
@@ -220,6 +221,32 @@ rules_py_test(
     deps = [":import_path_lib"],
 )
 
+# rules_python keeps `pyi_srcs` out of runfiles and in PyInfo only. Across the
+# ruleset boundary rules_py must pick them up from the foreign provider and
+# lay them down beside the modules in the venv.
+rules_python_library(
+    name = "stubbed_lib",
+    srcs = ["stubbed/stubbed.py"],
+    imports = ["stubbed"],
+    pyi_srcs = ["stubbed/stubbed.pyi"],
+)
+
+py_library(
+    name = "stubbed_wrapper",
+    deps = [":stubbed_lib"],
+)
+
+foreign_stubs_test(
+    name = "foreign_stubs_test",
+    target_under_test = ":stubbed_wrapper",
+)
+
+rules_py_test(
+    name = "rules_python_stubs_test",
+    srcs = ["rules_python_stubs_test.py"],
+    deps = [":stubbed_wrapper"],
+)
+
 # rules_python's console-script machinery reads the pip hub's dist-info to
 # generate a py_binary; unpinned, that binary lands on a rules_py-provisioned
 # toolchain. Its hub is parsed for 3.12 alone, so which wheel backs the script

e2e/rules-python-interop/rules_python_stubs_test.py

@@ -0,0 +1,14 @@
+"""A rules_python library's `pyi_srcs` reach a rules_py consumer's runfiles.
+
+The stub is metadata in rules_python's PyInfo, never in its own runfiles; the
+rules_py venv must carry it beside the module it annotates so a type checker
+pointed at the venv resolves it.
+"""
+
+import os
+
+import stubbed
+
+stub = os.path.splitext(stubbed.__file__)[0] + ".pyi"
+assert os.path.exists(stub), "missing type stub next to " + stubbed.__file__
+assert stubbed.describe(1) == "stubbed 1"

e2e/rules-python-interop/stubbed/stubbed.py

@@ -0,0 +1,2 @@
+def describe(value: object) -> str:
+    return "stubbed " + str(value)

e2e/rules-python-interop/stubbed/stubbed.pyi

@@ -0,0 +1 @@
+def describe(value: object) -> str: ...

e2e/rules-python-interop/type_stubs_test.bzl

@@ -0,0 +1,16 @@
+"""Analysis test: rules_python `pyi_srcs` surface in rules_py's `PyInfo`."""
+
+load("@aspect_rules_py//py:defs.bzl", "PyInfo")
+load("@bazel_skylib//lib:unittest.bzl", "analysistest", "asserts")
+
+def _basenames(files):
+    return sorted([file.basename for file in files.to_list()])
+
+def _foreign_stubs_test_impl(ctx):
+    env = analysistest.begin(ctx)
+    info = analysistest.target_under_test(env)[PyInfo]
+    asserts.equals(env, ["stubbed.pyi"], _basenames(info.transitive_pyi_files))
+    asserts.equals(env, ["stubbed.py"], _basenames(info.transitive_sources), "the stub is not a runtime source")
+    return analysistest.end(env)
+
+foreign_stubs_test = analysistest.make(_foreign_stubs_test_impl)

e2e/rules-python-protobuf/BUILD.bazel

@@ -1,6 +1,8 @@
-load("@aspect_rules_py//py:defs.bzl", "py_test")
+load("@aspect_rules_py//py:defs.bzl", "py_library", "py_test")
 load("@protobuf//bazel:proto_library.bzl", "proto_library")
+load("@protobuf//bazel:py_proto_library.bzl", "py_proto_library")
 load("@rules_proto_grpc_python//:defs.bzl", "python_proto_library")
+load(":type_stubs_test.bzl", "generated_stubs_test")
 
 # A plain proto_library — the language-agnostic descriptor. Living at the
 # module root keeps the generated module a clean top-level `greeting_pb2`.
@@ -28,3 +30,29 @@ py_test(
     main = "test.py",
     deps = [":greeting_py_proto"],
 )
+
+# protobuf's own py_proto_library runs protoc with `--pyi_out`, so the
+# generated bindings come with a `greeting_pb2.pyi` carried only in
+# rules_python's `PyInfo.transitive_pyi_files`.
+py_proto_library(
+    name = "greeting_py_proto_native",
+    deps = [":greeting_proto"],
+)
+
+py_library(
+    name = "greeting_stubs_wrapper",
+    deps = [":greeting_py_proto_native"],
+)
+
+generated_stubs_test(
+    name = "generated_stubs_test",
+    target_under_test = ":greeting_stubs_wrapper",
+)
+
+py_test(
+    name = "stubs_test",
+    size = "small",
+    srcs = ["stubs_test.py"],
+    main = "stubs_test.py",
+    deps = [":greeting_py_proto_native"],
+)

e2e/rules-python-protobuf/MODULE.bazel

@@ -13,6 +13,7 @@
 # own python_interpreters extension.
 
 bazel_dep(name = "aspect_rules_py")
+bazel_dep(name = "bazel_skylib", version = "1.4.2")
 bazel_dep(name = "rules_proto_grpc_python", version = "5.8.0")
 
 # Declared directly so the `@protobuf//bazel:proto_library.bzl` load in

e2e/rules-python-protobuf/stubs_test.py

@@ -0,0 +1,14 @@
+"""protobuf's py_proto_library emits `greeting_pb2.pyi` only through
+rules_python's `PyInfo.transitive_pyi_files`; the rules_py venv has to place it
+beside `greeting_pb2.py` for type checkers to see the generated message types.
+"""
+
+import os
+
+import greeting_pb2
+
+stub = os.path.splitext(greeting_pb2.__file__)[0] + ".pyi"
+assert os.path.exists(stub), "missing generated stub next to " + greeting_pb2.__file__
+
+with open(stub) as handle:
+    assert "class Greeting" in handle.read(), stub

e2e/rules-python-protobuf/type_stubs_test.bzl

@@ -0,0 +1,16 @@
+"""Analysis test: generated `_pb2.pyi` stubs surface in rules_py's `PyInfo`."""
+
+load("@aspect_rules_py//py:defs.bzl", "PyInfo")
+load("@bazel_skylib//lib:unittest.bzl", "analysistest", "asserts")
+
+def _basenames(files):
+    return sorted([file.basename for file in files.to_list()])
+
+def _generated_stubs_test_impl(ctx):
+    env = analysistest.begin(ctx)
+    info = analysistest.target_under_test(env)[PyInfo]
+    asserts.true(env, "greeting_pb2.pyi" in _basenames(info.transitive_pyi_files), "py_proto_library's stub propagates")
+    asserts.false(env, "greeting_pb2.pyi" in _basenames(info.transitive_sources), "the stub is not a runtime source")
+    return analysistest.end(env)
+
+generated_stubs_test = analysistest.make(_generated_stubs_test_impl)

e2e/rules-python-provider-compat/BUILD.bazel

@@ -1,5 +1,7 @@
 load("@aspect_rules_py//py:defs.bzl", "py_binary", "py_library", rules_py_test = "py_test")
+load("@rules_python//python:py_library.bzl", rules_python_library = "py_library")
 load("@rules_python//python:py_test.bzl", rules_python_test = "py_test")
+load(":type_stubs_test.bzl", "merged_stubs_test")
 
 # One converted library, two consumers: the @rules_python test stands in
 # for a package the migration hasn't reached yet, the rules_py test for one it
@@ -9,10 +11,27 @@ load("@rules_python//python:py_test.bzl", rules_python_test = "py_test")
 py_library(
     name = "lib",
     testonly = True,
-    srcs = ["lib.py"],
+    srcs = [
+        "lib.py",
+        "lib.pyi",
+    ],
     imports = ["."],
 )
 
+# rules_python's PyInfo builder merges the emitted provider field by field, so
+# the stub partition must arrive in the shape it expects: direct/transitive
+# `pyi_files` populated, and no `.pyi` among `transitive_sources`.
+rules_python_library(
+    name = "rules_python_wrapper",
+    testonly = True,
+    deps = [":lib"],
+)
+
+merged_stubs_test(
+    name = "merged_stubs_test",
+    target_under_test = ":rules_python_wrapper",
+)
+
 rules_python_test(
     name = "rules_python_consumer_test",
     srcs = ["consumer_test.py"],

e2e/rules-python-provider-compat/MODULE.bazel

@@ -11,6 +11,7 @@ requires.
 module(name = "rules_python_provider_compat")
 
 bazel_dep(name = "aspect_rules_py")
+bazel_dep(name = "bazel_skylib", version = "1.4.2")
 bazel_dep(name = "rules_python", version = "1.9.0")
 
 local_path_override(

e2e/rules-python-provider-compat/lib.pyi

@@ -0,0 +1 @@
+GREETING: str

e2e/rules-python-provider-compat/type_stubs_test.bzl

@@ -0,0 +1,17 @@
+"""Analysis test: a rules_py stub survives merging into rules_python's PyInfo."""
+
+load("@bazel_skylib//lib:unittest.bzl", "analysistest", "asserts")
+load("@rules_python//python:py_info.bzl", "PyInfo")
+
+def _basenames(files):
+    return sorted([file.basename for file in files.to_list()])
+
+def _merged_stubs_test_impl(ctx):
+    env = analysistest.begin(ctx)
+    info = analysistest.target_under_test(env)[PyInfo]
+    asserts.equals(env, [], _basenames(info.direct_pyi_files), "the wrapper declares no stubs of its own")
+    asserts.equals(env, ["lib.pyi"], _basenames(info.transitive_pyi_files))
+    asserts.equals(env, ["lib.py"], _basenames(info.transitive_sources), "the stub is not a runtime source")
+    return analysistest.end(env)
+
+merged_stubs_test = analysistest.make(_merged_stubs_test_impl)

py/private/py_info.bzl

@@ -11,7 +11,8 @@ deps to build the eventual venv or wheel.
 RulesPyInfo = provider(
     doc = "Python source, import-path, and virtual-dependency information for a target's dependency closure.",
     fields = {
-        "transitive_sources": "depset[File] — postorder depset of first-party `.py` sources in the transitive closure.",
+        "transitive_sources": "depset[File] — postorder depset of first-party runtime sources in the transitive closure; `.pyi` stubs are excluded.",
+        "transitive_pyi_files": "depset[File] — postorder depset of `.pyi` type stubs in the transitive closure: stubs listed in `srcs` plus those carried by deps of either ruleset.",
         "imports": "depset[str] — import roots to place on `sys.path` (rlocation-root-relative).",
         "virtual_dependencies": "depset[str] — names of required virtual dependencies, independent of their resolution status.",
         "virtual_resolutions": "depset[struct(virtual, target)] — virtual-dependency-name to concrete-target resolutions.",

py/private/py_info_interop.bzl

@@ -3,8 +3,8 @@
 The `deps` attribute on rules_py rules accepts targets built by either
 ruleset. rules_py always emits its own `PyInfo`
 (`//py/private:py_info.bzl`); native `@rules_python` targets (e.g.
-a `py_proto_library`) carry `@rules_python`'s. Both expose `transitive_sources`
-and `imports`, which is everything rules_py reads from a foreign dep.
+a `py_proto_library`) carry `@rules_python`'s. Both expose the source, type
+stub, and import information rules_py reads from a foreign dep.
 
 This module is the single place that knows about both providers. Rule code
 calls these accessors at the API edge instead of loading `@rules_python`'s
@@ -37,3 +37,11 @@ def get_py_info(target):
     if RulesPythonPyInfo in target:
         return target[RulesPythonPyInfo]
     return None
+
+def get_transitive_pyi_files(target):
+    """Return the `.pyi` closure from either ruleset's `PyInfo`, or an empty depset for targets carrying neither."""
+    if PyInfo in target:
+        return target[PyInfo].transitive_pyi_files
+    if RulesPythonPyInfo in target:
+        return target[RulesPythonPyInfo].transitive_pyi_files
+    return depset()