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 = select({ 26 "@platforms//os:windows": ["@platforms//:incompatible"], 27 "//conditions:default": [], 28}) if BZLMOD_ENABLED else ["@platforms//:incompatible"] 29 30def lock(*, name, srcs, out, upgrade = False, universal = True, python_version = None, args = [], **kwargs): 31 """Pin the requirements based on the src files. 32 33 Args: 34 name: The name of the target to run for updating the requirements. 35 srcs: The srcs to use as inputs. 36 out: The output file. 37 upgrade: Tell `uv` to always upgrade the dependencies instead of 38 keeping them as they are. 39 universal: Tell `uv` to generate a universal lock file. 40 python_version: Tell `rules_python` to use a particular version. 41 Defaults to the default py toolchain. 42 args: Extra args to pass to the rule. 43 **kwargs: Extra kwargs passed to the binary rule. 44 45 Differences with the current pip-compile rule: 46 - This is implemented in shell and uv. 47 - This does not error out if the output file does not exist yet. 48 - Supports transitions out of the box. 49 """ 50 pkg = native.package_name() 51 update_target = name + ".update" 52 53 _args = [ 54 "--custom-compile-command='bazel run //{}:{}'".format(pkg, update_target), 55 "--generate-hashes", 56 "--emit-index-url", 57 "--no-strip-extras", 58 "--python=$(PYTHON3)", 59 ] + args + [ 60 "$(location {})".format(src) 61 for src in srcs 62 ] 63 if upgrade: 64 _args.append("--upgrade") 65 if universal: 66 _args.append("--universal") 67 _args.append("--output-file=$@") 68 cmd = "$(UV_BIN) pip compile " + " ".join(_args) 69 70 # Make a copy to ensure that we are not modifying the initial list 71 srcs = list(srcs) 72 73 # Check if the output file already exists, if yes, first copy it to the 74 # output file location in order to make `uv` not change the requirements if 75 # we are just running the command. 76 if native.glob([out]): 77 cmd = "cp -v $(location {}) $@; {}".format(out, cmd) 78 srcs.append(out) 79 80 native.genrule( 81 name = name, 82 srcs = srcs, 83 outs = [out + ".new"], 84 cmd_bash = cmd, 85 tags = [ 86 "local", 87 "manual", 88 "no-cache", 89 ], 90 target_compatible_with = _REQUIREMENTS_TARGET_COMPATIBLE_WITH, 91 toolchains = [ 92 Label("//python/uv:current_toolchain"), 93 Label("//python:current_py_toolchain"), 94 ], 95 ) 96 if python_version: 97 py_binary_rule = lambda *args, **kwargs: transition_py_binary(python_version = python_version, *args, **kwargs) 98 else: 99 py_binary_rule = py_binary 100 101 # Write a script that can be used for updating the in-tree version of the 102 # requirements file 103 write_file( 104 name = name + ".update_gen", 105 out = update_target + ".py", 106 content = [ 107 "from os import environ", 108 "from pathlib import Path", 109 "from sys import stderr", 110 "", 111 'src = Path(environ["REQUIREMENTS_FILE"])', 112 'assert src.exists(), f"the {src} file does not exist"', 113 'dst = Path(environ["BUILD_WORKSPACE_DIRECTORY"]) / "{}" / "{}"'.format(pkg, out), 114 'print(f"Writing requirements contents\\n from {src.absolute()}\\n to {dst.absolute()}", file=stderr)', 115 "dst.write_text(src.read_text())", 116 'print("Success!", file=stderr)', 117 ], 118 ) 119 120 py_binary_rule( 121 name = update_target, 122 srcs = [update_target + ".py"], 123 main = update_target + ".py", 124 data = [name], 125 env = { 126 "REQUIREMENTS_FILE": "$(rootpath {})".format(name), 127 }, 128 tags = ["manual"], 129 **kwargs 130 ) 131