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