feat: entry_point support for bzlmod hub repos

aignas · aignas · commit a15a6e8da27a · 2023-07-07T14:20:46.000+09:00
diff --git a/examples/bzlmod/MODULE.bazel b/examples/bzlmod/MODULE.bazel
@@ -57,10 +57,7 @@ pip.parse(
     requirements_lock = "//:requirements_lock_3_10.txt",
     requirements_windows = "//:requirements_windows_3_10.txt",
 )
-
-# NOTE: The pip_39 repo is only used because the plain `@pip` repo doesn't
-# yet support entry points; see https://github.com/bazelbuild/rules_python/issues/1262
-use_repo(pip, "pip", "pip_39")
+use_repo(pip, "pip")
 
 bazel_dep(name = "other_module", version = "", repo_name = "our_other_module")
 local_path_override(
diff --git a/examples/bzlmod/entry_point/BUILD.bazel b/examples/bzlmod/entry_point/BUILD.bazel
@@ -1,4 +1,4 @@
-load("@pip_39//:requirements.bzl", "entry_point")
+load("@pip//:requirements.bzl", "entry_point")
 load("@rules_python//python:defs.bzl", "py_test")
 
 alias(
diff --git a/python/extensions/pip.bzl b/python/extensions/pip.bzl
@@ -15,7 +15,7 @@
 "pip module extension for use with bzlmod"
 
 load("@pythons_hub//:interpreters.bzl", "DEFAULT_PYTHON_VERSION", "INTERPRETER_LABELS")
-load("@rules_python//python/extensions/private:hub_repository.bzl", "hub_repository", "hub_repository_mega")
+load("@rules_python//python/extensions/private:hub_repository.bzl", "hub_repository")
 load(
     "@rules_python//python/pip_install:pip_repository.bzl",
     "locked_requirements_label",
@@ -24,6 +24,7 @@ load(
     "whl_library",
 )
 load("@rules_python//python/pip_install:requirements_parser.bzl", parse_requirements = "parse")
+load("//python:versions.bzl", "MINOR_MAPPING")
 
 def _create_versioned_pip_and_whl_repos(module_ctx, pip_attr, whl_map):
     python_interpreter_target = pip_attr.python_interpreter_target
@@ -55,20 +56,6 @@ def _create_versioned_pip_and_whl_repos(module_ctx, pip_attr, whl_map):
     requirements = parse_result.requirements
     extra_pip_args = pip_attr.extra_pip_args + parse_result.options
 
-    # Create the repository where users load the `requirement` macro. Under bzlmod
-    # this does not create the install_deps() macro.
-    # TODO: we may not need this repository once we have entry points
-    # supported. For now a user can access this repository and use
-    # the entrypoint functionality.
-    hub_repository(
-        name = pip_name,
-        repo_name = pip_name,
-        whl_library_alias_names = [
-            _sanitize_name(name)
-            for name, _ in parse_result.requirements
-        ],
-        template = "//python/extensions/private:hub_repository_default.bzl.tmpl",
-    )
     if hub_name not in whl_map:
         whl_map[hub_name] = {}
 
@@ -96,7 +83,7 @@ def _create_versioned_pip_and_whl_repos(module_ctx, pip_attr, whl_map):
         if whl_name not in whl_map[hub_name]:
             whl_map[hub_name][whl_name] = []
 
-        whl_map[hub_name][whl_name].append(pip_attr.python_version)
+        whl_map[hub_name][whl_name].append(MINOR_MAPPING.get(pip_attr.python_version, pip_attr.python_version))
 
 def _pip_impl(module_ctx):
     """Implementation of a class tag that creates the pip hub(s) and corresponding pip spoke, alias and whl repositories.
@@ -229,7 +216,7 @@ def _pip_impl(module_ctx):
 
     for hub_name, whl_map in hub_whl_map.items():
         # Create the hub repository for pip.
-        hub_repository_mega(
+        hub_repository(
             name = hub_name,
             repo_name = hub_name,
             whl_map = whl_map,
diff --git a/python/extensions/private/hub_repository.bzl b/python/extensions/private/hub_repository.bzl
@@ -14,6 +14,7 @@
 
 "A hub_repository rule used to create bzlmod hub repos."
 
+load("//python:versions.bzl", "MINOR_MAPPING")
 load("//python/private:create_pkg_aliases.bzl", "create_pkg_aliases")
 
 _BUILD_FILE_CONTENTS = """\
@@ -24,54 +25,10 @@ exports_files(["requirements.bzl"])
 """
 
 def _impl(rctx):
-    bzl_packages = rctx.attr.whl_library_alias_names
-    repo_name = rctx.attr.repo_name
-    build_contents = _BUILD_FILE_CONTENTS
-    create_pkg_aliases(rctx, repo_name, bzl_packages)
-
-    # NOTE: we are using the canonical name with the double '@' in order to
-    # always uniquely identify a repository, as the labels are being passed as
-    # a string and the resolution of the label happens at the call-site of the
-    # `requirement`, et al. macros.
-    macro_tmpl = "@@{name}//{{}}:{{}}".format(name = rctx.attr.name)
-
-    rctx.file("BUILD.bazel", build_contents)
-    rctx.template("requirements.bzl", rctx.attr.template, substitutions = {
-        "%%ALL_REQUIREMENTS%%": repr([
-            macro_tmpl.format(p, p)
-            for p in bzl_packages
-        ]),
-        "%%ALL_WHL_REQUIREMENTS%%": repr([
-            macro_tmpl.format(p, "whl")
-            for p in bzl_packages
-        ]),
-        "%%MACRO_TMPL%%": macro_tmpl,
-        "%%NAME%%": rctx.attr.name,
-    })
-
-hub_repository = repository_rule(
-    attrs = {
-        "repo_name": attr.string(
-            mandatory = True,
-            doc = "The apparent name of the repo. This is needed because in bzlmod, the name attribute becomes the canonical name.",
-        ),
-        "template": attr.label(
-            mandatory = True,
-            doc = "The template to use to create a hub repo",
-        ),
-        "whl_library_alias_names": attr.string_list(
-            mandatory = True,
-            doc = "The list of whl alias that we use to build aliases and the whl names",
-        ),
-    },
-    doc = """A rule for creating bzlmod hub repo. PRIVATE USE ONLY.""",
-    implementation = _impl,
-)
-
-def _impl_2(rctx):
     bzl_packages = rctx.attr.whl_map.keys()
     repo_name = rctx.attr.repo_name
     build_contents = _BUILD_FILE_CONTENTS
+
     create_pkg_aliases(
         rctx = rctx,
         repo_name = repo_name,
@@ -92,6 +49,7 @@ def _impl_2(rctx):
     # a string and the resolution of the label happens at the call-site of the
     # `requirement`, et al. macros.
     macro_tmpl = "@@{name}//{{}}:{{}}".format(name = rctx.attr.name)
+    entrypoint_tmpl = "@@{name}_{{version_label}}_{{pkg}}//:rules_python_wheel_entry_point_{{script}}".format(name = rctx.attr.name)
 
     rctx.file("BUILD.bazel", build_contents)
     rctx.template("requirements.bzl", rctx.attr._template, substitutions = {
@@ -103,11 +61,21 @@ def _impl_2(rctx):
             macro_tmpl.format(p, "whl")
             for p in bzl_packages
         ]),
+        "%%DEFAULT_PY_VERSION%%": repr(MINOR_MAPPING.get(rctx.attr.default_version, rctx.attr.default_version)),
+        "%%ENTRYPOINT_TMPL%%": entrypoint_tmpl,
         "%%MACRO_TMPL%%": macro_tmpl,
         "%%NAME%%": rctx.attr.name,
+        "%%PACKAGE_AVAILABILITY%%": repr({
+            k: [
+                MINOR_MAPPING.get(v, v)
+                for v in versions
+            ]
+            for k, versions in rctx.attr.whl_map.items()
+        }),
+        "%%RULES_PYTHON%%": rctx.attr._template.workspace_name,
     })
 
-hub_repository_mega = repository_rule(
+hub_repository = repository_rule(
     attrs = {
         "default_version": attr.string(
             mandatory = True,
@@ -118,13 +86,13 @@ hub_repository_mega = repository_rule(
         ),
         "whl_map": attr.string_list_dict(
             mandatory = True,
-            doc = "The wheel map where values are json strings encoding dicts.",
+            doc = "The wheel map where values are python versions",
         ),
         "_template": attr.label(
             doc = "The template to use to create a hub repo",
-            default = ":hub_repository_version.bzl.tmpl",
+            default = ":requirements.bzl.tmpl",
         ),
     },
     doc = """A rule for creating bzlmod hub repo. PRIVATE USE ONLY.""",
-    implementation = _impl_2,
+    implementation = _impl,
 )
diff --git a/python/extensions/private/hub_repository_default.bzl.tmpl b/python/extensions/private/hub_repository_default.bzl.tmpl
diff --git a/python/extensions/private/requirements.bzl.tmpl b/python/extensions/private/requirements.bzl.tmpl
@@ -6,10 +6,15 @@ This file is different from the other bzlmod template
 because we do not support entry_point yet.
 """
 
+load("@@%%RULES_PYTHON%%//python/pip_install:entrypoint.bzl", _entry_point = "entry_point")
+
 all_requirements = %%ALL_REQUIREMENTS%%
 
 all_whl_requirements = %%ALL_WHL_REQUIREMENTS%%
 
+_default_py_version = %%DEFAULT_PY_VERSION%%
+_packages = %%PACKAGE_AVAILABILITY%%
+
 def _clean_name(name):
     return name.replace("-", "_").replace(".", "_").lower()
 
@@ -28,5 +33,14 @@ def dist_info_requirement(name):
 def entry_point(pkg, script = None):
     """entry_point returns the target of the canonical label of the package entrypoints.
     """
-    # TODO: https://github.com/bazelbuild/rules_python/issues/1262
-    print("not implemented")
+    selects = _entry_point(
+        tmpl = "%%ENTRYPOINT_TMPL%%",
+        pkg = _clean_name(pkg),
+        script = script,
+        packages = _packages,
+        default_version = _default_py_version,
+    )
+    if selects == None:
+        fail("Package '{}' does not exist, select one from: {}".format(pkg, _packages.keys()))
+
+    return select(selects)
diff --git a/python/pip_install/entrypoint.bzl b/python/pip_install/entrypoint.bzl
@@ -0,0 +1,58 @@
+# Copyright 2023 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.
+
+"entry_point macro implementation for bzlmod. PRIVATE ONLY."
+
+def entry_point(*, pkg, packages, default_version, tmpl, script = None):
+    """Return an entry_point script.
+
+    Args:
+        pkg: the package name.
+        script: the script name to use, defaults to pkg name.
+        packages: the dictionary of packages to python versions that are supported.
+        default_version: the default python version.
+        tmpl: the template that will be interpolated by this function. The
+            following keys are going to be replaced: 'version_label', 'pkg' and
+            'script'.
+
+    Returns:
+        A dict that can be used in select statement or None if the pkg is not
+        in the supplied packages dictionary.
+    """
+    if not script:
+        script = pkg
+
+    if pkg not in packages:
+        return None
+
+    selects = {}
+    default = ""
+    for full_version in packages[pkg]:
+        condition = str(Label("//python/config_settings:is_python_{}".format(full_version)))
+
+        entry_point = tmpl.format(
+            version_label = full_version.rpartition(".")[0].replace(".", ""),
+            pkg = pkg,
+            script = script,
+        )
+
+        if full_version == default_version:
+            default = entry_point
+        else:
+            selects[condition] = entry_point
+
+    if default:
+        selects["//conditions:default"] = default
+
+    return selects
diff --git a/python/pip_install/private/test/BUILD.bazel b/python/pip_install/private/test/BUILD.bazel
@@ -1,4 +1,5 @@
 load("@bazel_skylib//rules:diff_test.bzl", "diff_test")
+load(":entrypoint_test.bzl", "entry_point_tests")
 load(":requirements_parser_tests.bzl", parse_requirements_tests = "parse_tests")
 
 diff_test(
@@ -18,3 +19,5 @@ diff_test(
 )
 
 parse_requirements_tests(name = "test_parse_requirements")
+
+entry_point_tests(name = "test_entry_point")
diff --git a/python/pip_install/private/test/entrypoint_test.bzl b/python/pip_install/private/test/entrypoint_test.bzl
diff --git a/python/private/create_pkg_aliases.bzl b/python/private/create_pkg_aliases.bzl

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-load("@pip_39//:requirements.bzl", "entry_point")`
	`1`	`+load("@pip//:requirements.bzl", "entry_point")`
`2`	`2`	`load("@rules_python//python:defs.bzl", "py_test")`
`3`	`3`
`4`	`4`	`alias(`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`load("@bazel_skylib//rules:diff_test.bzl", "diff_test")`
	`2`	`+load(":entrypoint_test.bzl", "entry_point_tests")`
`2`	`3`	`load(":requirements_parser_tests.bzl", parse_requirements_tests = "parse_tests")`
`3`	`4`
`4`	`5`	`diff_test(`
`@@ -18,3 +19,5 @@ diff_test(`
`18`	`19`	`)`
`19`	`20`
`20`	`21`	`parse_requirements_tests(name = "test_parse_requirements")`
	`22`	`+`
	`23`	`+entry_point_tests(name = "test_entry_point")`