From 261f2e3b5da5bd64e0baf09b875d2430feb7c96d Mon Sep 17 00:00:00 2001
From: Chris Chua <chua@uber.com>
Date: Sat, 17 Aug 2024 09:11:26 +0800
Subject: [PATCH] Enable passing in custom py_runtime (#98)

Could use some help and opinions on the attribute name.
This does add rules_python as a bazel dep in order to use the
`PyRuntimeInfo` provider. Not sure if that's undesirable. It seems to
make sense since rules_uv already depends on the toolchain made
available by it.

Example usage:
```starlark
# in MODULE.bazel
python = use_extension("@rules_python//python/extensions:python.bzl", "python")
python.toolchain(
    python_version = "3.11",
)
use_repo(
    python,
    "python_3_11",
)

# in BUILD.bazel
pip_compile(
    name = "requirements_3_11",
    py3_runtime = "@python_3_11//:py3_runtime",
    requirements_in = "requirements.in",
    requirements_txt = "requirements_3_11.txt",
)
```

Fixes #97
---
 MODULE.bazel       |  1 +
 uv/private/pip.bzl | 27 ++++++++++++++++++---------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/MODULE.bazel b/MODULE.bazel
index a943f91..10c6461 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -10,6 +10,7 @@ bazel_dep(name = "bazel_skylib", version = "1.4.1")
 bazel_dep(name = "buildifier_prebuilt", version = "6.1.2")
 bazel_dep(name = "platforms", version = "0.0.8")
 bazel_dep(name = "rules_multitool", version = "0.4.0")
+bazel_dep(name = "rules_python", version = "0.34.0")
 
 multitool = use_extension("@rules_multitool//multitool:extension.bzl", "multitool")
 multitool.hub(lockfile = "//uv/private:uv.lock.json")
diff --git a/uv/private/pip.bzl b/uv/private/pip.bzl
index e55460c..1dd6335 100644
--- a/uv/private/pip.bzl
+++ b/uv/private/pip.bzl
@@ -1,5 +1,7 @@
 "uv based pip compile rules"
 
+load("@rules_python//python:defs.bzl", "PyRuntimeInfo")
+
 _PY_TOOLCHAIN = "@bazel_tools//tools/python:toolchain_type"
 
 _DEFAULT_ARGS = [
@@ -12,34 +14,41 @@ _COMMON_ATTRS = {
     "requirements_in": attr.label(mandatory = True, allow_single_file = True),
     "requirements_txt": attr.label(mandatory = True, allow_single_file = True),
     "python_platform": attr.string(),
+    "py3_runtime": attr.label(),
     "data": attr.label_list(allow_files = True),
     "uv_args": attr.string_list(default = _DEFAULT_ARGS),
     "_uv": attr.label(default = "@multitool//tools/uv", executable = True, cfg = "exec"),
 }
 
-def _python_version(py_toolchain):
+def _python_version(py3_runtime):
     # micro is useful when there are some packages that exclude .0 versions due
     # to a bug, such as !=3.11.0
     return "{major}.{minor}.{micro}".format(
-        major = py_toolchain.py3_runtime.interpreter_version_info.major,
-        minor = py_toolchain.py3_runtime.interpreter_version_info.minor,
-        micro = py_toolchain.py3_runtime.interpreter_version_info.micro,
+        major = py3_runtime.interpreter_version_info.major,
+        minor = py3_runtime.interpreter_version_info.minor,
+        micro = py3_runtime.interpreter_version_info.micro,
     )
 
+def _python_runtime(ctx):
+    if ctx.attr.py3_runtime:
+        return ctx.attr.py3_runtime[PyRuntimeInfo]
+    py_toolchain = ctx.toolchains[_PY_TOOLCHAIN]
+    return py_toolchain.py3_runtime
+
 def _uv_pip_compile(
         ctx,
         template,
         executable,
         generator_label,
         uv_args):
-    py_toolchain = ctx.toolchains[_PY_TOOLCHAIN]
+    py3_runtime = _python_runtime(ctx)
     compile_command = "bazel run {label}".format(label = str(generator_label))
 
     args = []
     args += uv_args
     args.append("--custom-compile-command='{compile_command}'".format(compile_command = compile_command))
-    args.append("--python={python}".format(python = py_toolchain.py3_runtime.interpreter.short_path))
-    args.append("--python-version={version}".format(version = _python_version(py_toolchain)))
+    args.append("--python={python}".format(python = py3_runtime.interpreter.short_path))
+    args.append("--python-version={version}".format(version = _python_version(py3_runtime)))
     if ctx.attr.python_platform:
         args.append("--python-platform={platform}".format(platform = ctx.attr.python_platform))
 
@@ -56,10 +65,10 @@ def _uv_pip_compile(
     )
 
 def _runfiles(ctx):
-    py_toolchain = ctx.toolchains[_PY_TOOLCHAIN]
+    py3_runtime = _python_runtime(ctx)
     runfiles = ctx.runfiles(
         files = [ctx.file.requirements_in, ctx.file.requirements_txt] + ctx.files.data,
-        transitive_files = py_toolchain.py3_runtime.files,
+        transitive_files = py3_runtime.files,
     )
     runfiles = runfiles.merge(ctx.attr._uv.default_runfiles)
     return runfiles