feat: add pyi attributes/fields, original source fields (#2538)

This adds attributes and fields of use to static analysis.

For type definition files (usually `.pyi` files), the `pyi_srcs` and
`pyi_deps`
fields are added to the rules. They end up in the PyInfo fields
direct_pyi_files
and transitive_pyi_files.

So that static analysis tools can retain access to a target's Python
source files,
even if precompiling is enabled, `direct_original_sources` and
`transitive_original_sources` fields are added to PyInfo.

Work towards #2537,
#296

Author: rickeylev

CHANGELOG.md

@@ -97,6 +97,13 @@ Unreleased changes template.
     * 3.11.11
     * 3.12.8
     * 3.13.1
+* (rules) Attributes for type definition files (`.pyi` files) and type-checking
+  only dependencies added. See {obj}`py_library.pyi_srcs` and
+  `py_library.pyi_deps` (and the same named attributes for `py_binary` and
+  `py_test`).
+* (providers) {obj}`PyInfo` has new fields to aid static analysis tools:
+  {obj}`direct_original_sources`, {obj}`direct_pyi_files`,
+  {obj}`transitive_original_sources`, {obj}`transitive_pyi_files`.
 
 [20241206]: https://github.com/astral-sh/python-build-standalone/releases/tag/20241206
 

python/private/attributes.bzl

@@ -375,6 +375,35 @@ in the resulting output or not. Valid values are:
 * `omit_source`: Don't include the original py source.
 """,
         ),
+        "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,
+        ),
+        "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,
+        ),
         # Required attribute, but details vary by rule.
         # Use create_srcs_attr to create one.
         "srcs": None,

python/private/common.bzl

@@ -408,6 +408,7 @@ def collect_runfiles(ctx, files = depset()):
 def create_py_info(
         ctx,
         *,
+        original_sources,
         required_py_files,
         required_pyc_files,
         implicit_pyc_files,
@@ -417,6 +418,7 @@ def create_py_info(
 
     Args:
         ctx: rule ctx.
+        original_sources: `depset[File]`; the original input sources from `srcs`
         required_py_files: `depset[File]`; the direct, `.py` sources for the
             target that **must** be included by downstream targets. This should
             only be Python source files. It should not include pyc files.
@@ -435,10 +437,13 @@ def create_py_info(
         transitive sources collected from dependencies (the latter is only
         necessary for deprecated extra actions support).
     """
-
     py_info = PyInfoBuilder()
+    py_info.direct_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_original_sources.add(original_sources)
     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)
@@ -457,6 +462,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)

python/private/py_executable.bzl

@@ -985,6 +985,7 @@ def py_executable_base_impl(ctx, *, semantics, is_test, inherited_environment =
         runfiles_details = runfiles_details,
         main_py = main_py,
         imports = imports,
+        original_sources = direct_sources,
         required_py_files = required_py_files,
         required_pyc_files = required_pyc_files,
         implicit_pyc_files = implicit_pyc_files,
@@ -1548,6 +1549,7 @@ def _create_providers(
         ctx,
         executable,
         main_py,
+        original_sources,
         required_py_files,
         required_pyc_files,
         implicit_pyc_files,
@@ -1566,6 +1568,8 @@ def _create_providers(
         ctx: The rule ctx.
         executable: File; the target's executable file.
         main_py: File; the main .py entry point.
+        original_sources: `depset[File]` the direct `.py` sources for the
+            target that were the original input sources.
         required_py_files: `depset[File]` the direct, `.py` sources for the
             target that **must** be included by downstream targets. This should
             only be Python source files. It should not include pyc files.
@@ -1649,6 +1653,7 @@ def _create_providers(
 
     py_info, deps_transitive_sources, builtin_py_info = create_py_info(
         ctx,
+        original_sources = original_sources,
         required_py_files = required_py_files,
         required_pyc_files = required_pyc_files,
         implicit_pyc_files = implicit_pyc_files,

python/private/py_info.bzl

@@ -38,7 +38,11 @@ def _PyInfo_init(
         direct_pyc_files = depset(),
         transitive_pyc_files = depset(),
         transitive_implicit_pyc_files = depset(),
-        transitive_implicit_pyc_source_files = depset()):
+        transitive_implicit_pyc_source_files = depset(),
+        direct_original_sources = depset(),
+        transitive_original_sources = depset(),
+        direct_pyi_files = depset(),
+        transitive_pyi_files = depset()):
     _check_arg_type("transitive_sources", "depset", transitive_sources)
 
     # Verify it's postorder compatible, but retain is original ordering.
@@ -53,14 +57,24 @@ def _PyInfo_init(
 
     _check_arg_type("transitive_implicit_pyc_files", "depset", transitive_pyc_files)
     _check_arg_type("transitive_implicit_pyc_source_files", "depset", transitive_pyc_files)
+
+    _check_arg_type("direct_original_sources", "depset", direct_original_sources)
+    _check_arg_type("transitive_original_sources", "depset", transitive_original_sources)
+
+    _check_arg_type("direct_pyi_files", "depset", direct_pyi_files)
+    _check_arg_type("transitive_pyi_files", "depset", transitive_pyi_files)
     return {
+        "direct_original_sources": direct_original_sources,
         "direct_pyc_files": direct_pyc_files,
+        "direct_pyi_files": direct_pyi_files,
         "has_py2_only_sources": has_py2_only_sources,
         "has_py3_only_sources": has_py2_only_sources,
         "imports": imports,
         "transitive_implicit_pyc_files": transitive_implicit_pyc_files,
         "transitive_implicit_pyc_source_files": transitive_implicit_pyc_source_files,
+        "transitive_original_sources": transitive_original_sources,
         "transitive_pyc_files": transitive_pyc_files,
+        "transitive_pyi_files": transitive_pyi_files,
         "transitive_sources": transitive_sources,
         "uses_shared_libraries": uses_shared_libraries,
     }
@@ -69,6 +83,18 @@ PyInfo, _unused_raw_py_info_ctor = define_bazel_6_provider(
     doc = "Encapsulates information provided by the Python rules.",
     init = _PyInfo_init,
     fields = {
+        "direct_original_sources": """
+:type: depset[File]
+
+The `.py` source files (if any) that are considered directly provided by
+the target. This field is intended so that static analysis tools can recover the
+original Python source files, regardless of any build settings (e.g.
+precompiling), so they can analyze source code. The values are typically the
+`.py` files in the `srcs` attribute (or equivalent).
+
+::::{versionadded} 1.1.0
+::::
+""",
         "direct_pyc_files": """
 :type: depset[File]
 
@@ -78,6 +104,21 @@ by the target and **must be included**.
 These files usually come from, e.g., a library setting {attr}`precompile=enabled`
 to forcibly enable precompiling for itself. Downstream binaries are expected
 to always include these files, as the originating target expects them to exist.
+""",
+        "direct_pyi_files": """
+:type: depset[File]
+
+Type definition files (usually `.pyi` files) for the Python modules provided by
+this target. Usually they describe the source files listed in
+`direct_original_sources`. This field is primarily for static analysis tools.
+
+:::{note}
+This may contain implementation-specific file types specific to a particular
+type checker.
+:::
+
+::::{versionadded} 1.1.0
+::::
 """,
         "has_py2_only_sources": """
 :type: bool
@@ -116,6 +157,21 @@ then {obj}`transitive_implicit_pyc_files` should be included instead.
 
 ::::{versionadded} 0.37.0
 ::::
+""",
+        "transitive_original_sources": """
+:type: depset[File]
+
+The transitive set of `.py` source files (if any) that are considered the
+original sources for this target and its transitive dependencies. This field is
+intended so that static analysis tools can recover the original Python source
+files, regardless of any build settings (e.g. precompiling), so they can analyze
+source code. The values are typically the `.py` files in the `srcs` attribute
+(or equivalent).
+
+This is superset of `direct_original_sources`.
+
+::::{versionadded} 1.1.0
+::::
 """,
         "transitive_pyc_files": """
 :type: depset[File]
@@ -125,6 +181,22 @@ The transitive set of precompiled files that must be included.
 These files usually come from, e.g., a library setting {attr}`precompile=enabled`
 to forcibly enable precompiling for itself. Downstream binaries are expected
 to always include these files, as the originating target expects them to exist.
+""",
+        "transitive_pyi_files": """
+:type: depset[File]
+
+The transitive set of type definition files (usually `.pyi` files) for the
+Python modules for this target and its transitive dependencies. this target.
+Usually they describe the source files listed in `transitive_original_sources`.
+This field is primarily for static analysis tools.
+
+:::{note}
+This may contain implementation-specific file types specific to a particular
+type checker.
+:::
+
+::::{versionadded} 1.1.0
+::::
 """,
         "transitive_sources": """\
 :type: depset[File]
@@ -165,7 +237,9 @@ def PyInfoBuilder():
         _uses_shared_libraries = [False],
         build = lambda *a, **k: _PyInfoBuilder_build(self, *a, **k),
         build_builtin_py_info = lambda *a, **k: _PyInfoBuilder_build_builtin_py_info(self, *a, **k),
+        direct_original_sources = builders.DepsetBuilder(),
         direct_pyc_files = builders.DepsetBuilder(),
+        direct_pyi_files = builders.DepsetBuilder(),
         get_has_py2_only_sources = lambda *a, **k: _PyInfoBuilder_get_has_py2_only_sources(self, *a, **k),
         get_has_py3_only_sources = lambda *a, **k: _PyInfoBuilder_get_has_py3_only_sources(self, *a, **k),
         get_uses_shared_libraries = lambda *a, **k: _PyInfoBuilder_get_uses_shared_libraries(self, *a, **k),
@@ -182,7 +256,9 @@ def PyInfoBuilder():
         set_uses_shared_libraries = lambda *a, **k: _PyInfoBuilder_set_uses_shared_libraries(self, *a, **k),
         transitive_implicit_pyc_files = builders.DepsetBuilder(),
         transitive_implicit_pyc_source_files = builders.DepsetBuilder(),
+        transitive_original_sources = builders.DepsetBuilder(),
         transitive_pyc_files = builders.DepsetBuilder(),
+        transitive_pyi_files = builders.DepsetBuilder(),
         transitive_sources = builders.DepsetBuilder(),
     )
     return self
@@ -221,13 +297,39 @@ def _PyInfoBuilder_set_uses_shared_libraries(self, value):
     return self
 
 def _PyInfoBuilder_merge(self, *infos, direct = []):
+    """Merge other PyInfos into this PyInfo.
+
+    Args:
+        self: implicitly added.
+        *infos: {type}`PyInfo` objects to merge in, but only merge in their
+            information into this object's transitive fields.
+        direct: {type}`list[PyInfo]` objects to merge in, but also merge their
+            direct fields into this object's direct fields.
+
+    Returns:
+        {type}`PyInfoBuilder` the current object
+    """
     return self.merge_all(list(infos), direct = direct)
 
 def _PyInfoBuilder_merge_all(self, transitive, *, direct = []):
+    """Merge other PyInfos into this PyInfo.
+
+    Args:
+        self: implicitly added.
+        transitive: {type}`list[PyInfo]` objects to merge in, but only merge in
+            their information into this object's transitive fields.
+        direct: {type}`list[PyInfo]` objects to merge in, but also merge their
+            direct fields into this object's direct fields.
+
+    Returns:
+        {type}`PyInfoBuilder` the current object
+    """
     for info in direct:
         # BuiltinPyInfo doesn't have this field
         if hasattr(info, "direct_pyc_files"):
+            self.direct_original_sources.add(info.direct_original_sources)
             self.direct_pyc_files.add(info.direct_pyc_files)
+            self.direct_pyi_files.add(info.direct_pyi_files)
 
     for info in direct + transitive:
         self.imports.add(info.imports)
@@ -240,29 +342,58 @@ def _PyInfoBuilder_merge_all(self, transitive, *, direct = []):
         if hasattr(info, "transitive_pyc_files"):
             self.transitive_implicit_pyc_files.add(info.transitive_implicit_pyc_files)
             self.transitive_implicit_pyc_source_files.add(info.transitive_implicit_pyc_source_files)
+            self.transitive_original_sources.add(info.transitive_original_sources)
             self.transitive_pyc_files.add(info.transitive_pyc_files)
+            self.transitive_pyi_files.add(info.transitive_pyi_files)
 
     return self
 
 def _PyInfoBuilder_merge_target(self, target):
+    """Merge a target's Python information in this object.
+
+    Args:
+        self: implicitly added.
+        target: {type}`Target` targets that provide PyInfo, or other relevant
+        providers, will be merged into this object. If a target doesn't provide
+        any relevant providers, it is ignored.
+
+    Returns:
+        {type}`PyInfoBuilder` the current object.
+    """
     if PyInfo in target:
         self.merge(target[PyInfo])
     elif BuiltinPyInfo != None and BuiltinPyInfo in target:
         self.merge(target[BuiltinPyInfo])
     return self
 
 def _PyInfoBuilder_merge_targets(self, targets):
+    """Merge multiple targets into this object.
+
+    Args:
+        self: implicitly added.
+        targets: {type}`list[Target]`
+        targets that provide PyInfo, or other relevant
+        providers, will be merged into this object. If a target doesn't provide
+        any relevant providers, it is ignored.
+
+    Returns:
+        {type}`PyInfoBuilder` the current object.
+    """
     for t in targets:
         self.merge_target(t)
     return self
 
 def _PyInfoBuilder_build(self):
     if config.enable_pystar:
         kwargs = dict(
+            direct_original_sources = self.direct_original_sources.build(),
             direct_pyc_files = self.direct_pyc_files.build(),
-            transitive_pyc_files = self.transitive_pyc_files.build(),
+            direct_pyi_files = self.direct_pyi_files.build(),
             transitive_implicit_pyc_files = self.transitive_implicit_pyc_files.build(),
             transitive_implicit_pyc_source_files = self.transitive_implicit_pyc_source_files.build(),
+            transitive_original_sources = self.transitive_original_sources.build(),
+            transitive_pyc_files = self.transitive_pyc_files.build(),
+            transitive_pyi_files = self.transitive_pyi_files.build(),
         )
     else:
         kwargs = {}

python/private/py_library.bzl

@@ -102,6 +102,7 @@ def py_library_impl(ctx, *, semantics):
     cc_info = semantics.get_cc_info_for_library(ctx)
     py_info, deps_transitive_sources, builtins_py_info = create_py_info(
         ctx,
+        original_sources = direct_sources,
         required_py_files = required_py_files,
         required_pyc_files = required_pyc_files,
         implicit_pyc_files = implicit_pyc_files,