xref: /aosp_15_r20/external/bazelbuild-rules_python/python/uv/private/lock.bzl (revision 60517a1edbc8ecf509223e9af94a7adec7d736b8)
1# Copyright 2024 The Bazel Authors. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#    http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""A simple macro to lock the requirements.
16"""
17
18load("@bazel_skylib//rules:write_file.bzl", "write_file")
19load("//python:py_binary.bzl", "py_binary")
20load("//python/config_settings:transition.bzl", transition_py_binary = "py_binary")
21load("//python/private:bzlmod_enabled.bzl", "BZLMOD_ENABLED")  # buildifier: disable=bzl-visibility
22
23visibility(["//..."])
24
25_REQUIREMENTS_TARGET_COMPATIBLE_WITH = [] if BZLMOD_ENABLED else ["@platforms//:incompatible"]
26
27def lock(*, name, srcs, out, upgrade = False, universal = True, python_version = None):
28    """Pin the requirements based on the src files.
29
30    Args:
31        name: The name of the target to run for updating the requirements.
32        srcs: The srcs to use as inputs.
33        out: The output file.
34        upgrade: Tell `uv` to always upgrade the dependencies instead of
35            keeping them as they are.
36        universal: Tell `uv` to generate a universal lock file.
37        python_version: Tell `rules_python` to use a particular version.
38            Defaults to the default py toolchain.
39
40    Differences with the current pip-compile rule:
41    - This is implemented in shell and uv.
42    - This does not error out if the output file does not exist yet.
43    - Supports transitions out of the box.
44    """
45    pkg = native.package_name()
46    update_target = name + ".update"
47
48    args = [
49        "--custom-compile-command='bazel run //{}:{}'".format(pkg, update_target),
50        "--generate-hashes",
51        "--emit-index-url",
52        "--no-strip-extras",
53        "--python=$(PYTHON3)",
54    ] + [
55        "$(location {})".format(src)
56        for src in srcs
57    ]
58    if upgrade:
59        args.append("--upgrade")
60    if universal:
61        args.append("--universal")
62    args.append("--output-file=$@")
63    cmd = "$(UV_BIN) pip compile " + " ".join(args)
64
65    # Check if the output file already exists, if yes, first copy it to the
66    # output file location in order to make `uv` not change the requirements if
67    # we are just running the command.
68    if native.glob([out]):
69        cmd = "cp -v $(location {}) $@; {}".format(out, cmd)
70        srcs.append(out)
71
72    native.genrule(
73        name = name,
74        srcs = srcs,
75        outs = [out + ".new"],
76        cmd_bash = cmd,
77        tags = [
78            "local",
79            "manual",
80            "no-cache",
81        ],
82        target_compatible_with = _REQUIREMENTS_TARGET_COMPATIBLE_WITH,
83        toolchains = [
84            Label("//python/uv:current_toolchain"),
85            Label("//python:current_py_toolchain"),
86        ],
87    )
88    if python_version:
89        py_binary_rule = lambda *args, **kwargs: transition_py_binary(python_version = python_version, *args, **kwargs)
90    else:
91        py_binary_rule = py_binary
92
93    # Write a script that can be used for updating the in-tree version of the
94    # requirements file
95    write_file(
96        name = name + ".update_gen",
97        out = update_target + ".py",
98        content = [
99            "from os import environ",
100            "from pathlib import Path",
101            "from sys import stderr",
102            "",
103            'src = Path(environ["REQUIREMENTS_FILE"])',
104            'assert src.exists(), f"the {src} file does not exist"',
105            'dst = Path(environ["BUILD_WORKSPACE_DIRECTORY"]) / "{}" / "{}"'.format(pkg, out),
106            'print(f"Writing requirements contents\\n  from {src.absolute()}\\n  to {dst.absolute()}", file=stderr)',
107            "dst.write_text(src.read_text())",
108            'print("Success!", file=stderr)',
109        ],
110    )
111
112    py_binary_rule(
113        name = update_target,
114        srcs = [update_target + ".py"],
115        main = update_target + ".py",
116        data = [name],
117        env = {
118            "REQUIREMENTS_FILE": "$(rootpath {})".format(name),
119        },
120        tags = ["manual"],
121    )
122