feat: add api for analysis-phase code

CHANGELOG.md

@@ -33,6 +33,8 @@ A brief description of the categories of changes:
 ### Added
 * (py_wheel) Now supports `compress = (True|False)` to allow disabling
   compression to speed up development.
+* (api) Added {obj}`merge_py_infos()` so user rules can merge and propagate
+  `PyInfo` without losing information.
 
 ### Removed
 * Nothing yet

docs/BUILD.bazel

@@ -93,10 +93,12 @@ sphinx_stardocs(
         "//python:py_runtime_info_bzl",
         "//python:py_test_bzl",
         "//python:repositories_bzl",
+        "//python/api:api_bzl",
         "//python/cc:py_cc_toolchain_bzl",
         "//python/cc:py_cc_toolchain_info_bzl",
         "//python/entry_points:py_console_script_binary_bzl",
         "//python/private:py_cc_toolchain_rule_bzl",
+        "//python/private/api:py_common_api_bzl",
         "//python/private/common:py_binary_rule_bazel_bzl",
         "//python/private/common:py_library_rule_bazel_bzl",
         "//python/private/common:py_runtime_rule_bzl",

python/BUILD.bazel

@@ -34,6 +34,7 @@ licenses(["notice"])
 filegroup(
     name = "distribution",
     srcs = glob(["**"]) + [
+        "//python/api:distribution",
         "//python/cc:distribution",
         "//python/config_settings:distribution",
         "//python/constraints:distribution",

python/api/BUILD.bazel

@@ -0,0 +1,31 @@
+# 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.
+
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+
+package(
+    default_visibility = ["//:__subpackages__"],
+)
+
+bzl_library(
+    name = "api_bzl",
+    srcs = ["api.bzl"],
+    visibility = ["//visibility:public"],
+    deps = ["//python/private/api:api_bzl"],
+)
+
+filegroup(
+    name = "distribution",
+    srcs = glob(["**"]),
+)

python/api/api.bzl

@@ -0,0 +1,5 @@
+"""Public, analysis phase APIs for Python rules."""
+
+load("//python/private/api:api.bzl", _py_common = "py_common")
+
+py_common = _py_common

python/private/api/BUILD.bazel

@@ -0,0 +1,43 @@
+# 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.
+
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+load(":py_common_api.bzl", "py_common_api")
+
+package(
+    default_visibility = ["//:__subpackages__"],
+)
+
+py_common_api(
+    name = "py_common_api",
+    # NOTE: Not actually public. Implicit dependency of public rules.
+    visibility = ["//visibility:public"],
+)
+
+bzl_library(
+    name = "api_bzl",
+    srcs = ["api.bzl"],
+    deps = [
+        "//python/private:py_info_bzl",
+    ],
+)
+
+bzl_library(
+    name = "py_common_api_bzl",
+    srcs = ["py_common_api.bzl"],
+    deps = [
+        ":api_bzl",
+        "//python/private:py_info_bzl",
+    ],
+)

python/private/api/api.bzl

@@ -0,0 +1,55 @@
+# 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.
+"""Implementation of py_api."""
+
+_PY_COMMON_API_LABEL = Label("//python/private/api:py_common_api")
+
+ApiImplInfo = provider(
+    doc = "Provider to hold an API implementation",
+    fields = {
+        "impl": """
+:type: struct
+
+The implementation of the API being provided. The object it contains
+will depend on the target that is providing the API struct.
+""",
+    },
+)
+
+def _py_common_get(ctx):
+    """Get the py_common API instance.
+
+    NOTE: to use this function, the rule must have added `py_common.API_ATTRS`
+    to its attributes.
+
+    Args:
+        ctx: {type}`ctx` current rule ctx
+
+    Returns:
+        {type}`PyCommonApi`
+    """
+
+    # A generic provider is used to decouple the API implementations from
+    # the loading phase of the rules using an implementation.
+    return ctx.attr._py_common_api[ApiImplInfo].impl
+
+py_common = struct(
+    get = _py_common_get,
+    API_ATTRS = {
+        "_py_common_api": attr.label(
+            default = _PY_COMMON_API_LABEL,
+            providers = [ApiImplInfo],
+        ),
+    },
+)

python/private/api/py_common_api.bzl

@@ -0,0 +1,38 @@
+# 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.
+"""Implementation of py_api."""
+
+load("//python/private:py_info.bzl", "PyInfoBuilder")
+load("//python/private/api:api.bzl", "ApiImplInfo")
+
+def _py_common_api_impl(ctx):
+    _ = ctx  # @unused
+    return [ApiImplInfo(impl = PyCommonApi)]
+
+py_common_api = rule(
+    implementation = _py_common_api_impl,
+    doc = "Rule implementing py_common API.",
+)
+
+def _merge_py_infos(transitive, *, direct = []):
+    builder = PyInfoBuilder()
+    builder.merge_all(transitive, direct = direct)
+    return builder.build()
+
+# Exposed for doc generation, not directly used.
+# buildifier: disable=name-conventions
+PyCommonApi = struct(
+    merge_py_infos = _merge_py_infos,
+    PyInfoBuilder = PyInfoBuilder,
+)

python/private/py_info.bzl

@@ -181,11 +181,16 @@ def _PyInfoBuilder_set_uses_shared_libraries(self, value):
     self._uses_shared_libraries[0] = value
     return self
 
-def _PyInfoBuilder_merge(self, *infos):
-    return self.merge_all(infos)
+def _PyInfoBuilder_merge(self, *infos, direct = []):
+    return self.merge_all(list(infos), direct = direct)
 
-def _PyInfoBuilder_merge_all(self, py_infos):
-    for info in py_infos:
+def _PyInfoBuilder_merge_all(self, transitive, *, direct = []):
+    for info in direct:
+        # BuiltinPyInfo doesn't have this field
+        if hasattr(info, "direct_pyc_files"):
+            self.direct_pyc_files.add(info.direct_pyc_files)
+
+    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)

tests/api/py_common/BUILD.bazel

@@ -0,0 +1,17 @@
+# 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.
+
+load(":py_common_tests.bzl", "py_common_test_suite")
+
+py_common_test_suite(name = "py_common_tests")

tests/api/py_common/py_common_tests.bzl

@@ -0,0 +1,64 @@
+# 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.
+"""py_common tests."""
+
+load("@rules_testing//lib:analysis_test.bzl", "analysis_test")
+load("@rules_testing//lib:test_suite.bzl", "test_suite")
+load("@rules_testing//lib:util.bzl", rt_util = "util")
+load("//python/api:api.bzl", _py_common = "py_common")
+load("//tests/support:py_info_subject.bzl", "py_info_subject")
+
+_tests = []
+
+def _test_merge_py_infos(name):
+    rt_util.helper_target(
+        native.filegroup,
+        name = name + "_subject",
+        srcs = ["f1.py", "f1.pyc", "f2.py", "f2.pyc"],
+    )
+    analysis_test(
+        name = name,
+        impl = _test_merge_py_infos_impl,
+        target = name + "_subject",
+        attrs = _py_common.API_ATTRS,
+    )
+
+def _test_merge_py_infos_impl(env, target):
+    f1_py, f1_pyc, f2_py, f2_pyc = target[DefaultInfo].files.to_list()
+
+    py_common = _py_common.get(env.ctx)
+
+    py1 = py_common.PyInfoBuilder()
+    py1.direct_pyc_files.add(f1_pyc)
+    py1.transitive_sources.add(f1_py)
+
+    py2 = py_common.PyInfoBuilder()
+    py1.direct_pyc_files.add(f2_pyc)
+    py2.transitive_sources.add(f2_py)
+
+    actual = py_info_subject(
+        py_common.merge_py_infos([py2.build()], direct = [py1.build()]),
+        meta = env.expect.meta,
+    )
+
+    actual.transitive_sources().contains_exactly([f1_py.path, f2_py.path])
+    actual.direct_pyc_files().contains_exactly([f1_pyc.path, f2_pyc.path])
+
+_tests.append(_test_merge_py_infos)
+
+def py_common_test_suite(name):
+    test_suite(
+        name = name,
+        tests = _tests,
+    )