feat: add public API for analysis-phase logic (#2252)

This adds a public API for rules (i.e. analysis-phase code) to use code
from rules_python.
The main motivation for this is so that users can propagate PyInfo
without having to know
all the fields of PyInfo and implement the merging logic. With upcoming
PRs adding additional
fields to PyInfo, this becomes much more important.

The way the API is exposed is through a target. There are three reasons
for this:
1. It avoids loading phase costs when the implementation of the API
functions change.
Within Google, this makes changes to rules_python much cheaper and
easier to submit
and revert. This also allows us to worry less about the loading-phase
impact of
   our code.
2. Because a target can have dependencies, it allows us to hide some
details
from users. For example, if we want a flag to affect behavior, we can
add it to the
API target's attributes; users don't have to add it to their rule's
attributes
3. By having the API take the user's `ctx` as an argument, it allows us
to capture it
and use it as part of future API calls (this isn't used now, but gives
us
   flexibility in the future).

Work towards https://github.com/bazelbuild/rules_python/issues/1647
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3077756..d58351a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -47,6 +47,8 @@
 * (toolchains): A public `//python/config_settings:python_version_major_minor` has
   been exposed for users to be able to match on the `X.Y` version of a Python
   interpreter.
+* (api) Added {obj}`merge_py_infos()` so user rules can merge and propagate
+  `PyInfo` without losing information.
 
 ### Removed
 * Nothing yet
diff --git a/docs/BUILD.bazel b/docs/BUILD.bazel
index 149e2c5..66b6496 100644
--- a/docs/BUILD.bazel
+++ b/docs/BUILD.bazel
@@ -93,10 +93,12 @@
         "//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",
diff --git a/python/BUILD.bazel b/python/BUILD.bazel
index 53fb812..e64ad8c 100644
--- a/python/BUILD.bazel
+++ b/python/BUILD.bazel
@@ -34,6 +34,7 @@
 filegroup(
     name = "distribution",
     srcs = glob(["**"]) + [
+        "//python/api:distribution",
         "//python/cc:distribution",
         "//python/config_settings:distribution",
         "//python/constraints:distribution",
diff --git a/python/api/BUILD.bazel b/python/api/BUILD.bazel
new file mode 100644
index 0000000..1df6877
--- /dev/null
+++ b/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(["**"]),
+)
diff --git a/python/api/api.bzl b/python/api/api.bzl
new file mode 100644
index 0000000..c8fb921
--- /dev/null
+++ b/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
diff --git a/python/private/api/BUILD.bazel b/python/private/api/BUILD.bazel
new file mode 100644
index 0000000..9e97dc2
--- /dev/null
+++ b/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",
+    ],
+)
diff --git a/python/private/api/api.bzl b/python/private/api/api.bzl
new file mode 100644
index 0000000..06fb729
--- /dev/null
+++ b/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],
+        ),
+    },
+)
diff --git a/python/private/api/py_common_api.bzl b/python/private/api/py_common_api.bzl
new file mode 100644
index 0000000..401b359
--- /dev/null
+++ b/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,
+)
diff --git a/python/private/py_info.bzl b/python/private/py_info.bzl
index a3e40f2..97cd50b 100644
--- a/python/private/py_info.bzl
+++ b/python/private/py_info.bzl
@@ -181,11 +181,16 @@
     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)
diff --git a/tests/api/py_common/BUILD.bazel b/tests/api/py_common/BUILD.bazel
new file mode 100644
index 0000000..0930037
--- /dev/null
+++ b/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")
diff --git a/tests/api/py_common/py_common_tests.bzl b/tests/api/py_common/py_common_tests.bzl
new file mode 100644
index 0000000..572028b
--- /dev/null
+++ b/tests/api/py_common/py_common_tests.bzl
@@ -0,0 +1,68 @@
+# 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_python_internal//:rules_python_config.bzl", "config")
+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()
+    if config.enable_pystar:
+        py1.direct_pyc_files.add(f1_pyc)
+    py1.transitive_sources.add(f1_py)
+
+    py2 = py_common.PyInfoBuilder()
+    if config.enable_pystar:
+        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])
+    if config.enable_pystar:
+        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,
+    )
diff --git a/tests/base_rules/py_info/py_info_tests.bzl b/tests/base_rules/py_info/py_info_tests.bzl
index b64263f..97c8e26 100644
--- a/tests/base_rules/py_info/py_info_tests.bzl
+++ b/tests/base_rules/py_info/py_info_tests.bzl
@@ -111,30 +111,25 @@
         name = name + "_misc",
         srcs = ["trans.py", "direct.pyc", "trans.pyc"],
     )
-    rt_util.helper_target(
-        provide_py_info,
-        name = name + "_py1",
-        transitive_sources = ["py1-trans.py"],
-        direct_pyc_files = ["py1-direct-pyc.pyc"],
-        imports = ["py1import"],
-        transitive_pyc_files = ["py1-trans.pyc"],
-    )
-    rt_util.helper_target(
-        provide_py_info,
-        name = name + "_py2",
-        transitive_sources = ["py2-trans.py"],
-        direct_pyc_files = ["py2-direct.pyc"],
-        imports = ["py2import"],
-        transitive_pyc_files = ["py2-trans.pyc"],
-    )
+
+    py_info_targets = {}
+    for n in range(1, 7):
+        py_info_name = "{}_py{}".format(name, n)
+        py_info_targets["py{}".format(n)] = py_info_name
+        rt_util.helper_target(
+            provide_py_info,
+            name = py_info_name,
+            transitive_sources = ["py{}-trans.py".format(n)],
+            direct_pyc_files = ["py{}-direct.pyc".format(n)],
+            imports = ["py{}import".format(n)],
+            transitive_pyc_files = ["py{}-trans.pyc".format(n)],
+        )
     analysis_test(
         name = name,
         impl = _test_py_info_builder_impl,
         targets = {
             "misc": name + "_misc",
-            "py1": name + "_py1",
-            "py2": name + "_py2",
-        },
+        } | py_info_targets,
     )
 
 def _test_py_info_builder_impl(env, targets):
@@ -151,6 +146,9 @@
     builder.merge_target(targets.py1)
     builder.merge_targets([targets.py2])
 
+    builder.merge(targets.py3[PyInfo], direct = [targets.py4[PyInfo]])
+    builder.merge_all([targets.py5[PyInfo]], direct = [targets.py6[PyInfo]])
+
     def check(actual):
         subject = py_info_subject(actual, meta = env.expect.meta)
 
@@ -162,20 +160,34 @@
             "tests/base_rules/py_info/trans.py",
             "tests/base_rules/py_info/py1-trans.py",
             "tests/base_rules/py_info/py2-trans.py",
+            "tests/base_rules/py_info/py3-trans.py",
+            "tests/base_rules/py_info/py4-trans.py",
+            "tests/base_rules/py_info/py5-trans.py",
+            "tests/base_rules/py_info/py6-trans.py",
         ])
         subject.imports().contains_exactly([
             "import-path",
             "py1import",
             "py2import",
+            "py3import",
+            "py4import",
+            "py5import",
+            "py6import",
         ])
         if hasattr(actual, "direct_pyc_files"):
             subject.direct_pyc_files().contains_exactly([
                 "tests/base_rules/py_info/direct.pyc",
+                "tests/base_rules/py_info/py4-direct.pyc",
+                "tests/base_rules/py_info/py6-direct.pyc",
             ])
             subject.transitive_pyc_files().contains_exactly([
                 "tests/base_rules/py_info/trans.pyc",
                 "tests/base_rules/py_info/py1-trans.pyc",
                 "tests/base_rules/py_info/py2-trans.pyc",
+                "tests/base_rules/py_info/py3-trans.pyc",
+                "tests/base_rules/py_info/py4-trans.pyc",
+                "tests/base_rules/py_info/py5-trans.pyc",
+                "tests/base_rules/py_info/py6-trans.pyc",
             ])
 
     check(builder.build())