feat: add PyInfo fields for storing original sources of a target

Author: rickeylev

CHANGELOG.md

@@ -95,6 +95,10 @@ Unreleased changes template.
     * 3.11.11
     * 3.12.8
     * 3.13.1
+* (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`. NOTE: these
+  are not yet fully populated by `py_library`
 
 [20241206]: https://github.com/astral-sh/python-build-standalone/releases/tag/20241206
 

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.
@@ -437,6 +439,8 @@ def create_py_info(
     """
 
     py_info = PyInfoBuilder()
+    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.transitive_pyc_files.add(required_pyc_files)
     py_info.transitive_implicit_pyc_files.add(implicit_pyc_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,25 @@ 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_pyi_files": transitive_pyi_files,
         "transitive_sources": transitive_sources,
         "uses_shared_libraries": uses_shared_libraries,
     }
@@ -69,6 +84,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 +105,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 +158,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 +182,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 +238,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 +257,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
@@ -228,13 +305,15 @@ def _PyInfoBuilder_merge_all(self, transitive, *, direct = []):
         # BuiltinPyInfo doesn't have this field
         if hasattr(info, "direct_pyc_files"):
             self.direct_pyc_files.add(info.direct_pyc_files)
+            self.direct_original_sources.add(info.direct_original_sources)
 
     for info in direct + transitive:
         self.imports.add(info.imports)
         self.merge_has_py2_only_sources(info.has_py2_only_sources)
         self.merge_has_py3_only_sources(info.has_py3_only_sources)
         self.merge_uses_shared_libraries(info.uses_shared_libraries)
         self.transitive_sources.add(info.transitive_sources)
+        self.transitive_original_sources.add(info.transitive_original_sources)
 
         # BuiltinPyInfo doesn't have these fields
         if hasattr(info, "transitive_pyc_files"):
@@ -259,10 +338,14 @@ def _PyInfoBuilder_merge_targets(self, targets):
 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,