Initial crack at entrypoint support

Author: arrdem

e2e/MODULE.bazel

@@ -90,4 +90,9 @@ uv.lockfile(
     lockfile = "//cases/uv-deps-650:uv-airflow.lock",
     venv_name = "airflow",
 )
+
+# Ask uv to make entrypoints for setuptools
+uv.discover_entrypoints(
+    requirement = "setuptools",
+)
 use_repo(uv, "pypi")

uv/private/extension.bzl

@@ -296,7 +296,8 @@ def _raw_whl_repos(module_ctx, lock_specs):
                     #
                     # Assume (potentially a problem!)
                     name = _whl_repo_name(package, whl)
-                    print("Creating whl repo", name)
+
+                    # print("Creating whl repo", name)
                     downloaded_file_path = url.split("/")[-1]
                     spec = dict(
                         name = name,
@@ -331,7 +332,8 @@ def _sbuild_repos(module_ctx, lock_specs):
                     continue
 
                 name = _sbuild_repo_name(hub_name, venv_name, package)
-                print("Creating sdist repo", name)
+
+                # print("Creating sdist repo", name)
                 sdist_build(
                     name = name,
                     src = "@" + _sdist_repo_name(package) + "//file",
@@ -348,6 +350,90 @@ def _whl_install_repo_name(hub, venv, package):
         package["name"],
     )
 
+def debug(it):
+    return "<%s %s>" % (type(it), {
+        k: getattr(it, k)
+        for k in dir(it)
+    })
+
+def parse_ini(lines):
+    dict = {}
+    heading = None
+    for line in lines.split("\n"):
+        line = line.strip()
+        if line.startswith("[") and line.endswith("]"):
+            heading = line[1:-2]
+            dict[heading] = {}
+
+        elif "=" in line and heading:
+            key, value = line.split("=", 1)
+            key = key.strip()
+            value = value.strip()
+            dict[heading][key] = value
+
+    return dict
+
+def _collect_entrypoints(module_ctx, lock_specs):
+    # Collect all the packages anyone's allowing entrypoints for
+    packages_to_collect = {}
+    for mod in module_ctx.modules:
+        for it in mod.tags.discover_entrypoints:
+            packages_to_collect[normalize_name(it.requirement)] = 1
+
+    entrypoints = {}
+
+    # Collect predeclared entrypoints
+    for mod in module_ctx.modules:
+        for it in mod.tags.declare_entrypoint:
+            r = normalize_name(it.requirement)
+            entrypoints.setdefault(r, {})
+            entrypoints[r][normalize_name(it.name)] = it.entrypoint
+
+    print(packages_to_collect)
+
+    for hub_name, venvs in lock_specs.items():
+        for venv_name, lock in venvs.items():
+            for package in lock.get("package", []):
+                # We use packages_to_collect as both a set and a worklist
+                if package["name"] not in packages_to_collect:
+                    continue
+
+                reference_prebuild = None
+
+                # Find the smallest available wheel
+                for whl in package.get("wheels", []):
+                    reference_prebuild = whl
+                    break
+
+                if not reference_prebuild:
+                    continue
+
+                whl_filename = whl["url"].split("/")[-1]
+                file = "private/" + whl_filename
+                module_ctx.download(reference_prebuild["url"], sha256 = reference_prebuild["hash"][len("sha256:"):], output = file)
+                res = module_ctx.execute(
+                    [
+                        "tar",  # FIXME: Use a hermetic tar here somehow?
+                        "-xOzf",
+                        file,
+                        "*.dist-info/entry_points.txt",
+                    ],
+                )
+                print(debug(res))
+                if res.return_code == 0:
+                    entrypoints.setdefault(package["name"], {})
+                    whl_entrypoints = parse_ini(res.stdout)
+                    print(package["name"], res.stdout, whl_entrypoints)
+                    for name, entrypoint in whl_entrypoints.get("console_script", {}).items():
+                        entrypoints[package["name"]][normalize_name(name)] = entrypoint
+
+                    packages_to_collect.pop(package["name"])
+
+                else:
+                    print(file, res.exit_code, res.stderr)
+
+    return entrypoints
+
 def _whl_install_repos(module_ctx, lock_specs):
     for hub_name, venvs in lock_specs.items():
         for venv_name, lock in venvs.items():
@@ -366,7 +452,8 @@ def _whl_install_repos(module_ctx, lock_specs):
                 # only with the single venv. Shouldn't be possible to force this
                 # target to build when the venv hub is not pointed to this venv.
                 name = _whl_install_repo_name(hub_name, venv_name, package)
-                print("Creating install repo", name)
+
+                # print("Creating install repo", name)
                 whl_install(
                     name = name,
                     prebuilds = json.encode(prebuilds),
@@ -385,7 +472,7 @@ def _marker_sha(marker):
     else:
         return None
 
-def _group_repos(module_ctx, lock_specs):
+def _group_repos(module_ctx, lock_specs, entrypoint_specs):
     # Hub -> requirement -> venv -> True
     # For building hubs we need to know what venv configurations a given
 
@@ -534,7 +621,8 @@ def _group_repos(module_ctx, lock_specs):
             # their direct dependencies beyond the scc. So we can just lay down
             # targets.
             name = _venv_hub_name(hub_name, venv_name)
-            print("Creating venv hub", name)
+
+            # print("Creating venv hub", name)
             venv_hub(
                 name = name,
                 aliases = scc_aliases,  # String dict
@@ -546,13 +634,14 @@ def _group_repos(module_ctx, lock_specs):
                     package: _whl_install_repo_name(hub_name, venv_name, {"name": package})
                     for package in sorted(graph.keys())
                 },
+                entrypoints = json.encode(entrypoint_specs),
             )
 
     return package_venvs
 
-def _hub_repos(module_ctx, lock_specs, package_venvs):
+def _hub_repos(module_ctx, lock_specs, package_venvs, entrypoint_specs):
     for hub_name, packages in package_venvs.items():
-        print("Creating uv hub", hub_name)
+        # print("Creating uv hub", hub_name)
         hub_repo(
             name = hub_name,
             hub_name = hub_name,
@@ -561,6 +650,7 @@ def _hub_repos(module_ctx, lock_specs, package_venvs):
                 package: venvs.keys()
                 for package, venvs in packages.items()
             },
+            entrypoints = json.encode(entrypoint_specs),
         )
 
 def _uv_impl(module_ctx):
@@ -578,6 +668,10 @@ def _uv_impl(module_ctx):
     # of conditions.
     configurations = _collect_configurations(module_ctx, lock_specs)
 
+    # Collect declared entrypoints for packages
+    entrypoints = _collect_entrypoints(module_ctx, lock_specs)
+    print("got entrypoints", entrypoints)
+
     # Roll through and create sdist and whl repos for all configured sources
     # Note that these have no deps to this point
     _raw_sdist_repos(module_ctx, lock_specs)
@@ -587,13 +681,18 @@ def _uv_impl(module_ctx):
     _sbuild_repos(module_ctx, lock_specs)
 
     # Roll through and create per-venv whl installs
+    #
+    # Note that we handle entrypoints at the venv level NOT the install level.
+    # This is because we handle cycle breaking and deps at the venv level, so we
+    # can't just take a direct dependency on the installed whl in its
+    # implementation repo.
     _whl_install_repos(module_ctx, lock_specs)
 
     # Roll through and create per-venv group/dep layers
-    package_venvs = _group_repos(module_ctx, lock_specs)
+    package_venvs = _group_repos(module_ctx, lock_specs, entrypoints)
 
     # Finally the hubs themselves are fully trivialized
-    _hub_repos(module_ctx, lock_specs, package_venvs)
+    _hub_repos(module_ctx, lock_specs, package_venvs, entrypoints)
 
     configurations_hub(
         name = "aspect_rules_py_pip_configurations",
@@ -622,11 +721,27 @@ _lockfile_tag = tag_class(
     },
 )
 
+_declare_entrypoint = tag_class(
+    attrs = {
+        "requirement": attr.string(mandatory = True),
+        "name": attr.string(mandatory = True),
+        "entrypoint": attr.string(mandatory = True),
+    },
+)
+
+_create_entrypoints = tag_class(
+    attrs = {
+        "requirement": attr.string(mandatory = True),
+    },
+)
+
 uv = module_extension(
     implementation = _uv_impl,
     tag_classes = {
         "declare_hub": _hub_tag,
         "declare_venv": _venv_tag,
         "lockfile": _lockfile_tag,
+        "declare_entrypoint": _declare_entrypoint,
+        "discover_entrypoints": _create_entrypoints,
     },
 )

uv/private/hub/repository.bzl

@@ -106,6 +106,8 @@ pip = struct(
     ################################################################################
     # Lay down the hub aliases
 
+    entrypoints = json.decode(repository_ctx.attr.entrypoints)
+
     # FIXME: since we're creating a package per target, we may have to implement
     # name mangling to ensure that the pip packages become valid Bazel packages.
     for name, spec in repository_ctx.attr.packages.items():
@@ -119,6 +121,7 @@ load("//:defs.bzl", "pip")
             "//venv:{}".format(it): "@venv__{0}__{1}//{2}:{2}".format(repository_ctx.attr.hub_name, it, name)
             for it in spec
         }
+        error = "Available only in venvs " + ", ".join([it.split(":")[1][1:] for it in select_spec.keys()])
 
         # TODO: Find a way to add a dist-info target here
         # TODO: Find a way to add entrypoint targets here?
@@ -140,18 +143,53 @@ alias(
                 name = name,
                 select = repr(select_spec),
                 compat = repr(spec),
-                error = "Available only in venvs " + ", ".join([it.split(":")[1][1:] for it in select_spec.keys()]),
+                error = error,
             ),
         )
 
         repository_ctx.file(name + "/BUILD.bazel", content = "\n".join(content))
 
+        content = [
+            """load("//:defs.bzl", "pip")""",
+        ]
+        for entrypoint_name, entrypoint_coordinate in entrypoints.get(name, {}).items():
+            select_spec = {
+                "//venv:{}".format(it): "@venv__{0}__{1}//{2}/entrypoints:{3}".format(repository_ctx.attr.hub_name, it, name, entrypoint_name)
+                for it in spec
+            }
+
+            content.append(
+                """
+alias(
+    name = "{name}",
+    actual = select(
+      {select},
+      no_match_error = "{error}",
+    ),
+    target_compatible_with = pip.compatible_with({compat}),
+    visibility = ["//visibility:public"],
+)
+""".format(
+                    name = entrypoint_name,
+                    select = repr(select_spec),
+                    compat = repr(spec),
+                    error = error,
+                ),
+            )
+
+        repository_ctx.file(name + "/entrypoints/BUILD.bazel", content = "\n".join(content))
+
 hub_repo = repository_rule(
     implementation = _hub_impl,
     attrs = {
         "hub_name": attr.string(),
         "venvs": attr.string_list(),
         "packages": attr.string_list_dict(),
         "configurations": attr.string_list_dict(),
+        "entrypoints": attr.string(
+            doc = """
+        JSON encoded map of pkg -> entrypoint -> coordinate
+        """,
+        ),
     },
 )

uv/private/py_entrypoint_binary/BUILD.bazel

@@ -0,0 +1 @@
+exports_files(["entrypoint.tmpl"])

uv/private/py_entrypoint_binary/defs.bzl

@@ -0,0 +1,40 @@
+load("@bazel_lib//lib:expand_template.bzl", "expand_template")
+load("//py/unstable:defs.bzl", "py_venv_binary")
+
+
+def py_entrypoint_binary(
+    name,
+    coordinate,
+    deps,
+    visibility = ["//visibility:public"],
+):
+    main = "_{}_entrypoint".format(name)
+    # <name> = <package_or_module>[:<object>[.<attr>[.<nested-attr>]*]]
+    package, symbol = coordinate.split(":")
+
+    if "." in symbol:
+        fn, tail = symbol.split(".", 1)
+        alias = "{fn} = {fn}.{tail}\n".format(fn=fn, tail=tail)
+    else:
+        fn = symbol
+        tail = ""
+        alias = ""
+
+    expand_template(
+        name = main,
+        template = Label("@aspect_rules_py//uv/private/py_entrypoint_binary:entrypoint.tmpl"),
+        out = main + ".py",
+        substitutions = {
+            "{{package}}": package,
+            "{{fn}}": fn,
+            "{{alias}}": alias
+        },
+    )
+    
+    py_venv_binary(
+        name = name,
+        main = main + ".py",
+        srcs = [main],
+        deps = deps,
+        visibility = visibility
+    )

uv/private/py_entrypoint_binary/entrypoint.tmpl

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+from {{package}} import {{fn}}
+{{alias}}
+if __name__ == '__main__':
+    {{fn}}()