feat: add pyi_srcs, pyi_deps attributes

Author: rickeylev

python/private/attributes.bzl

@@ -395,6 +395,35 @@ in the resulting output or not. Valid values are:
         # Force enabling auto exec groups, see
         # https://bazel.build/extending/auto-exec-groups#how-enable-particular-rule
         "_use_auto_exec_groups": attr.bool(default = True),
+        "pyi_srcs": attr.label_list(
+            doc = """
+Type definition files for the library.
+
+These are typically `.pyi` files, but other file types for type-checker specific
+formats are allowed. These files are build-time only dependencies and not included
+as part of a runnable program (packaging rules may include them, however).
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+            allow_files = True,
+        ),
+        "pyi_deps": attr.label_list(
+            doc = """
+Dependencies providing type definitions the library needs.
+
+These are dependencies that satisfy imports guarded by `typing.TYPE_CHECKING`.
+These are build-time only dependencies and not included as part of a runnable
+program (packaging rules may include them, however).
+
+:::{versionadded} VERSION_NEXT_FEATURE
+:::
+""",
+            providers = [
+                [PyInfo],
+                [CcInfo],
+            ] + _MaybeBuiltinPyInfo,
+        ),
     },
     allow_none = True,
 )

python/private/common.bzl

@@ -442,7 +442,9 @@ def create_py_info(
     py_info.direct_original_sources.add(original_sources)
     py_info.transitive_original_sources.add(original_sources)
     py_info.direct_pyc_files.add(required_pyc_files)
+    py_info.direct_pyi_files.add(ctx.files.pyi_srcs)
     py_info.transitive_pyc_files.add(required_pyc_files)
+    py_info.transitive_pyi_files.add(ctx.files.pyi_srcs)
     py_info.transitive_implicit_pyc_files.add(implicit_pyc_files)
     py_info.transitive_implicit_pyc_source_files.add(implicit_pyc_source_files)
     py_info.imports.add(imports)
@@ -461,6 +463,10 @@ def create_py_info(
                 if f.extension == "py":
                     py_info.transitive_sources.add(f)
                 py_info.merge_uses_shared_libraries(cc_helper.is_valid_shared_library_artifact(f))
+    for target in ctx.attr.pyi_deps:
+        # PyInfo may not be present e.g. cc_library rules.
+        if PyInfo in target or (BuiltinPyInfo != None and BuiltinPyInfo in target):
+            py_info.merge(_get_py_info(target))
 
     deps_transitive_sources = py_info.transitive_sources.build()
     py_info.transitive_sources.add(required_py_files)

tests/base_rules/base_tests.bzl

@@ -17,6 +17,7 @@ load("@rules_testing//lib:analysis_test.bzl", "analysis_test")
 load("@rules_testing//lib:truth.bzl", "matching")
 load("@rules_testing//lib:util.bzl", "PREVENT_IMPLICIT_BUILDING_TAGS", rt_util = "util")
 load("//python:py_info.bzl", "PyInfo")
+load("//python:py_library.bzl", "py_library")
 load("//python/private:reexports.bzl", "BuiltinPyInfo")  # buildifier: disable=bzl-visibility
 load("//tests/base_rules:util.bzl", pt_util = "util")
 load("//tests/support:py_info_subject.bzl", "py_info_subject")
@@ -58,6 +59,49 @@ _not_produces_py_info = rule(
     implementation = _not_produces_py_info_impl,
 )
 
+def _test_py_info_populated(name, config):
+    rt_util.helper_target(
+        config.base_test_rule,
+        name = name + "_subject",
+        pyi_srcs = ["subject.pyi"],
+        pyi_deps = [name + "_lib2"],
+    )
+    rt_util.helper_target(
+        py_library,
+        name = name + "_lib2",
+        srcs = ["lib2.py"],
+        pyi_srcs = ["lib2.pyi"],
+    )
+
+    analysis_test(
+        name = name,
+        target = name + "_subject",
+        impl = _test_py_info_populated_impl,
+    )
+
+def _test_py_info_populated_impl(env, target):
+    info = env.expect.that_target(target).provider(
+        PyInfo,
+        factory = py_info_subject,
+    )
+    info.direct_original_sources().contains_exactly([
+        "{package}/test_py_info_populated_subject.py",
+    ])
+    info.transitive_original_sources().contains_exactly([
+        "{package}/test_py_info_populated_subject.py",
+        "{package}/lib2.py",
+    ])
+
+    info.direct_pyi_files().contains_exactly([
+        "{package}/subject.pyi",
+    ])
+    info.transitive_pyi_files().contains_exactly([
+        "{package}/lib2.pyi",
+        "{package}/subject.pyi",
+    ])
+
+_tests.append(_test_py_info_populated)
+
 def _py_info_propagation_setup(name, config, produce_py_info_rule, test_impl):
     rt_util.helper_target(
         config.base_test_rule,