refactor: add builders to make creating PyInfo/depset/runfiles easier

python/private/BUILD.bazel

@@ -69,6 +69,14 @@ bzl_library(
     ],
 )
 
+bzl_library(
+    name = "builders_bzl",
+    srcs = ["builders.bzl"],
+    deps = [
+        "@bazel_skylib//lib:types",
+    ],
+)
+
 bzl_library(
     name = "bzlmod_enabled_bzl",
     srcs = ["bzlmod_enabled.bzl"],
@@ -270,8 +278,10 @@ bzl_library(
     name = "py_info_bzl",
     srcs = ["py_info.bzl"],
     deps = [
+        ":builders_bzl",
         ":reexports_bzl",
         ":util_bzl",
+        "@rules_python_internal//:rules_python_config_bzl",
     ],
 )
 

python/private/builders.bzl

@@ -0,0 +1,190 @@
+# Copyright 2024 The Bazel Authors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Builders to make building complex objects easier."""
+
+load("@bazel_skylib//lib:types.bzl", "types")
+
+def _DepsetBuilder():
+    """Create a builder for a depset."""
+
+    # buildifier: disable=uninitialized
+    self = struct(
+        _order = [None],
+        add = lambda *a, **k: _DepsetBuilder_add(self, *a, **k),
+        build = lambda *a, **k: _DepsetBuilder_build(self, *a, **k),
+        direct = [],
+        get_order = lambda *a, **k: _DepsetBuilder_get_order(self, *a, **k),
+        set_order = lambda *a, **k: _DepsetBuilder_set_order(self, *a, **k),
+        transitive = [],
+    )
+    return self
+
+def _DepsetBuilder_add(self, *values):
+    """Add value to the depset.
+
+    Args:
+        self: {type}`DepsetBuilder` implicitly added.
+        *values: {type}`depset | list | object` Values to add to the depset.
+            The values can be a depset, the non-depset value to add, or
+            a list of such values to add.
+
+    Returns:
+        {type}`DepsetBuilder`
+    """
+    for value in values:
+        if types.is_list(value):
+            for sub_value in value:
+                if types.is_depset(sub_value):
+                    self.transitive.append(sub_value)
+                else:
+                    self.direct.append(sub_value)
+        elif types.is_depset(value):
+            self.transitive.append(value)
+        else:
+            self.direct.append(value)
+    return self
+
+def _DepsetBuilder_set_order(self, order):
+    """Sets the order to use.
+
+    Args:
+        self: {type}`DepsetBuilder` implicitly added.
+        order: {type}`str` One of the {obj}`depset` `order` values.
+
+    Returns:
+        {type}`DepsetBuilder`
+    """
+    self._order[0] = order
+    return self
+
+def _DepsetBuilder_get_order(self):
+    """Gets the depset order that will be used.
+
+    Args:
+        self: {type}`DepsetBuilder` implicitly added.
+
+    Returns:
+        {type}`str | None` If not previously set, `None` is returned.
+    """
+    return self._order[0]
+
+def _DepsetBuilder_build(self):
+    """Creates a {obj}`depset` from the accumulated values.
+
+    Args:
+        self: {type}`DepsetBuilder` implicitly added.
+
+    Returns:
+        {type}`depset`
+    """
+    if not self.direct and len(self.transitive) == 1 and self._order[0] == None:
+        return self.transitive[0]
+    else:
+        kwargs = {}
+        if self._order[0] != None:
+            kwargs["order"] = self._order[0]
+        return depset(direct = self.direct, transitive = self.transitive, **kwargs)
+
+def _RunfilesBuilder():
+    """Creates a `RunfilesBuilder`.
+
+    Returns:
+        {type}`RunfilesBuilder`
+    """
+
+    # buildifier: disable=uninitialized
+    self = struct(
+        add = lambda *a, **k: _RunfilesBuilder_add(self, *a, **k),
+        add_targets = lambda *a, **k: _RunfilesBuilder_add_targets(self, *a, **k),
+        build = lambda *a, **k: _RunfilesBuilder_build(self, *a, **k),
+        files = _DepsetBuilder(),
+        root_symlinks = {},
+        runfiles = [],
+        symlinks = {},
+    )
+    return self
+
+def _RunfilesBuilder_add(self, *values):
+    """Adds a value to the runfiles.
+
+    Args:
+        self: {type}`RunfilesBuilder` implicitly added.
+        *values: {type}`File | runfiles | list[File] | depset[File] | list[runfiles]`
+            The values to add.
+
+    Returns:
+        {type}`RunfilesBuilder`
+    """
+    for value in values:
+        if types.is_list(value):
+            for sub_value in value:
+                _RunfilesBuilder_add_internal(self, sub_value)
+        else:
+            _RunfilesBuilder_add_internal(self, value)
+    return self
+
+def _RunfilesBuilder_add_targets(self, targets):
+    """Adds runfiles from targets
+
+    Args:
+        self: {type}`RunfilesBuilder` implicitly added.
+        targets: {type}`list[Target]` targets whose default runfiles
+            to add.
+
+    Returns:
+        {type}`RunfilesBuilder`
+    """
+    for t in targets:
+        self.runfiles.append(t[DefaultInfo].default_runfiles)
+    return self
+
+def _RunfilesBuilder_add_internal(self, value):
+    if _is_file(value):
+        self.files.add(value)
+    elif types.is_depset(value):
+        self.files.add(value)
+    elif _is_runfiles(value):
+        self.runfiles.append(value)
+    else:
+        fail("Unhandled value: type {}: {}".format(type(value), value))
+
+def _RunfilesBuilder_build(self, ctx, **kwargs):
+    """Creates a {obj}`runfiles` from the accumulated values.
+
+    Args:
+        self: {type}`RunfilesBuilder` implicitly added.
+        ctx: {type}`ctx` The rule context to use to create the runfiles object.
+        **kwargs: additional args to pass along to {obj}`ctx.runfiles`.
+
+    Returns:
+        {type}`runfiles`
+    """
+    return ctx.runfiles(
+        transitive_files = self.files.build(),
+        symlinks = self.symlinks,
+        root_symlinks = self.root_symlinks,
+        **kwargs
+    ).merge_all(self.runfiles)
+
+# Skylib's types module doesn't have is_file, so roll our own
+def _is_file(value):
+    return type(value) == "File"
+
+def _is_runfiles(value):
+    return type(value) == "runfiles"
+
+builders = struct(
+    DepsetBuilder = _DepsetBuilder,
+    RunfilesBuilder = _RunfilesBuilder,
+)

python/private/common/common.bzl

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Various things common to Bazel and Google rule implementations."""
 
-load("//python/private:py_info.bzl", "PyInfo")
+load("//python/private:py_info.bzl", "PyInfo", "PyInfoBuilder")
 load("//python/private:reexports.bzl", "BuiltinPyInfo")
 load(":cc_helper.bzl", "cc_helper")
 load(":py_internal.bzl", "py_internal")
@@ -282,7 +282,7 @@ def collect_imports(ctx, semantics):
         if BuiltinPyInfo in dep
     ])
 
-def collect_runfiles(ctx, files):
+def collect_runfiles(ctx, files = depset()):
     """Collects the necessary files from the rule's context.
 
     This presumes the ctx is for a py_binary, py_test, or py_library rule.
@@ -364,84 +364,50 @@ def create_py_info(ctx, *, direct_sources, direct_pyc_files, imports):
         transitive sources collected from dependencies (the latter is only
         necessary for deprecated extra actions support).
     """
-    uses_shared_libraries = False
-    has_py2_only_sources = ctx.attr.srcs_version in ("PY2", "PY2ONLY")
-    has_py3_only_sources = ctx.attr.srcs_version in ("PY3", "PY3ONLY")
-    transitive_sources_depsets = []  # list of depsets
-    transitive_sources_files = []  # list of Files
-    transitive_pyc_depsets = [direct_pyc_files]  # list of depsets
+
+    py_info = PyInfoBuilder()
+    py_info.direct_pyc_files.add(direct_pyc_files)
+    py_info.transitive_pyc_files.add(direct_pyc_files)
+    py_info.imports.add(imports)
+    py_info.merge_has_py2_only_sources(ctx.attr.srcs_version in ("PY2", "PY2ONLY"))
+    py_info.merge_has_py3_only_sources(ctx.attr.srcs_version in ("PY3", "PY3ONLY"))
+
     for target in ctx.attr.deps:
         # PyInfo may not be present e.g. cc_library rules.
         if PyInfo in target or BuiltinPyInfo in target:
-            info = _get_py_info(target)
-            transitive_sources_depsets.append(info.transitive_sources)
-            uses_shared_libraries = uses_shared_libraries or info.uses_shared_libraries
-            has_py2_only_sources = has_py2_only_sources or info.has_py2_only_sources
-            has_py3_only_sources = has_py3_only_sources or info.has_py3_only_sources
-
-            # BuiltinPyInfo doesn't have this field.
-            if hasattr(info, "transitive_pyc_files"):
-                transitive_pyc_depsets.append(info.transitive_pyc_files)
+            py_info.merge(_get_py_info(target))
         else:
             # TODO(b/228692666): Remove this once non-PyInfo targets are no
             # longer supported in `deps`.
             files = target.files.to_list()
             for f in files:
                 if f.extension == "py":
-                    transitive_sources_files.append(f)
-                uses_shared_libraries = (
-                    uses_shared_libraries or
-                    cc_helper.is_valid_shared_library_artifact(f)
-                )
-    deps_transitive_sources = depset(
-        direct = transitive_sources_files,
-        transitive = transitive_sources_depsets,
-    )
+                    py_info.transitive_sources.add(f)
+                py_info.merge_uses_shared_libraries(cc_helper.is_valid_shared_library_artifact(f))
+
+    deps_transitive_sources = py_info.transitive_sources.build()
+    py_info.transitive_sources.add(direct_sources)
 
     # We only look at data to calculate uses_shared_libraries, if it's already
     # true, then we don't need to waste time looping over it.
-    if not uses_shared_libraries:
+    if not py_info.get_uses_shared_libraries():
         # Similar to the above, except we only calculate uses_shared_libraries
         for target in ctx.attr.data:
             # TODO(b/234730058): Remove checking for PyInfo in data once depot
             # cleaned up.
             if PyInfo in target or BuiltinPyInfo in target:
                 info = _get_py_info(target)
-                uses_shared_libraries = info.uses_shared_libraries
+                py_info.merge_uses_shared_libraries(info.uses_shared_libraries)
             else:
                 files = target.files.to_list()
                 for f in files:
-                    uses_shared_libraries = cc_helper.is_valid_shared_library_artifact(f)
-                    if uses_shared_libraries:
+                    py_info.merge_uses_shared_libraries(cc_helper.is_valid_shared_library_artifact(f))
+                    if py_info.get_uses_shared_libraries():
                         break
-            if uses_shared_libraries:
+            if py_info.get_uses_shared_libraries():
                 break
 
-    py_info_kwargs = dict(
-        transitive_sources = depset(
-            transitive = [deps_transitive_sources, direct_sources],
-        ),
-        imports = imports,
-        # NOTE: This isn't strictly correct, but with Python 2 gone,
-        # the srcs_version logic is largely defunct, so shouldn't matter in
-        # practice.
-        has_py2_only_sources = has_py2_only_sources,
-        has_py3_only_sources = has_py3_only_sources,
-        uses_shared_libraries = uses_shared_libraries,
-        direct_pyc_files = direct_pyc_files,
-        transitive_pyc_files = depset(transitive = transitive_pyc_depsets),
-    )
-
-    # TODO(b/203567235): Set `uses_shared_libraries` field, though the Bazel
-    # docs indicate it's unused in Bazel and may be removed.
-    py_info = PyInfo(**py_info_kwargs)
-
-    # Remove args that BuiltinPyInfo doesn't support
-    py_info_kwargs.pop("direct_pyc_files")
-    py_info_kwargs.pop("transitive_pyc_files")
-    builtin_py_info = BuiltinPyInfo(**py_info_kwargs)
-
-    return py_info, deps_transitive_sources, builtin_py_info
+    return py_info.build(), deps_transitive_sources, py_info.build_builtin_py_info()
 
 def _get_py_info(target):
     return target[PyInfo] if PyInfo in target else target[BuiltinPyInfo]