1# Copyright 2023 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"Set defaults for the pip-compile command to run it under Bazel" 16 17import atexit 18import os 19import shutil 20import sys 21from pathlib import Path 22 23import piptools.writer as piptools_writer 24from piptools.scripts.compile import cli 25 26from python.runfiles import runfiles 27 28# Replace the os.replace function with shutil.copy to work around os.replace not being able to 29# replace or move files across filesystems. 30os.replace = shutil.copy 31 32# Next, we override the annotation_style_split and annotation_style_line functions to replace the 33# backslashes in the paths with forward slashes. This is so that we can have the same requirements 34# file on Windows and Unix-like. 35original_annotation_style_split = piptools_writer.annotation_style_split 36original_annotation_style_line = piptools_writer.annotation_style_line 37 38 39def annotation_style_split(required_by) -> str: 40 required_by = set([v.replace("\\", "/") for v in required_by]) 41 return original_annotation_style_split(required_by) 42 43 44def annotation_style_line(required_by) -> str: 45 required_by = set([v.replace("\\", "/") for v in required_by]) 46 return original_annotation_style_line(required_by) 47 48 49piptools_writer.annotation_style_split = annotation_style_split 50piptools_writer.annotation_style_line = annotation_style_line 51 52 53def _select_golden_requirements_file( 54 requirements_txt, requirements_linux, requirements_darwin, requirements_windows 55): 56 """Switch the golden requirements file, used to validate if updates are needed, 57 to a specified platform specific one. Fallback on the platform independent one. 58 """ 59 60 plat = sys.platform 61 if plat == "linux" and requirements_linux is not None: 62 return requirements_linux 63 elif plat == "darwin" and requirements_darwin is not None: 64 return requirements_darwin 65 elif plat == "win32" and requirements_windows is not None: 66 return requirements_windows 67 else: 68 return requirements_txt 69 70 71def _locate(bazel_runfiles, file): 72 """Look up the file via Rlocation""" 73 74 if not file: 75 return file 76 77 return bazel_runfiles.Rlocation(file) 78 79 80if __name__ == "__main__": 81 if len(sys.argv) < 4: 82 print( 83 "Expected at least two arguments: requirements_in requirements_out", 84 file=sys.stderr, 85 ) 86 sys.exit(1) 87 88 parse_str_none = lambda s: None if s == "None" else s 89 bazel_runfiles = runfiles.Create() 90 91 requirements_in = sys.argv.pop(1) 92 requirements_txt = sys.argv.pop(1) 93 requirements_linux = parse_str_none(sys.argv.pop(1)) 94 requirements_darwin = parse_str_none(sys.argv.pop(1)) 95 requirements_windows = parse_str_none(sys.argv.pop(1)) 96 update_target_label = sys.argv.pop(1) 97 98 requirements_file = _select_golden_requirements_file( 99 requirements_txt=requirements_txt, requirements_linux=requirements_linux, 100 requirements_darwin=requirements_darwin, requirements_windows=requirements_windows 101 ) 102 103 resolved_requirements_in = _locate(bazel_runfiles, requirements_in) 104 resolved_requirements_file = _locate(bazel_runfiles, requirements_file) 105 106 # Files in the runfiles directory has the following naming schema: 107 # Main repo: __main__/<path_to_file> 108 # External repo: <workspace name>/<path_to_file> 109 # We want to strip both __main__ and <workspace name> from the absolute prefix 110 # to keep the requirements lock file agnostic. 111 repository_prefix = requirements_file[: requirements_file.index("/") + 1] 112 absolute_path_prefix = resolved_requirements_file[ 113 : -(len(requirements_file) - len(repository_prefix)) 114 ] 115 116 # As requirements_in might contain references to generated files we want to 117 # use the runfiles file first. Thus, we need to compute the relative path 118 # from the execution root. 119 # Note: Windows cannot reference generated files without runfiles support enabled. 120 requirements_in_relative = requirements_in[len(repository_prefix):] 121 requirements_file_relative = requirements_file[len(repository_prefix):] 122 123 # Before loading click, set the locale for its parser. 124 # If it leaks through to the system setting, it may fail: 125 # RuntimeError: Click will abort further execution because Python 3 was configured to use ASCII 126 # as encoding for the environment. Consult https://click.palletsprojects.com/python3/ for 127 # mitigation steps. 128 os.environ["LC_ALL"] = "C.UTF-8" 129 os.environ["LANG"] = "C.UTF-8" 130 131 UPDATE = True 132 # Detect if we are running under `bazel test`. 133 if "TEST_TMPDIR" in os.environ: 134 UPDATE = False 135 # pip-compile wants the cache files to be writeable, but if we point 136 # to the real user cache, Bazel sandboxing makes the file read-only 137 # and we fail. 138 # In theory this makes the test more hermetic as well. 139 sys.argv.append("--cache-dir") 140 sys.argv.append(os.environ["TEST_TMPDIR"]) 141 # Make a copy for pip-compile to read and mutate. 142 requirements_out = os.path.join( 143 os.environ["TEST_TMPDIR"], os.path.basename(requirements_file) + ".out" 144 ) 145 # Those two files won't necessarily be on the same filesystem, so we can't use os.replace 146 # or shutil.copyfile, as they will fail with OSError: [Errno 18] Invalid cross-device link. 147 shutil.copy(resolved_requirements_file, requirements_out) 148 149 update_command = os.getenv("CUSTOM_COMPILE_COMMAND") or "bazel run %s" % ( 150 update_target_label, 151 ) 152 153 os.environ["CUSTOM_COMPILE_COMMAND"] = update_command 154 os.environ["PIP_CONFIG_FILE"] = os.getenv("PIP_CONFIG_FILE") or os.devnull 155 156 sys.argv.append("--output-file") 157 sys.argv.append(requirements_file_relative if UPDATE else requirements_out) 158 sys.argv.append( 159 requirements_in_relative 160 if Path(requirements_in_relative).exists() 161 else resolved_requirements_in 162 ) 163 print(sys.argv) 164 165 if UPDATE: 166 print("Updating " + requirements_file_relative) 167 if "BUILD_WORKSPACE_DIRECTORY" in os.environ: 168 workspace = os.environ["BUILD_WORKSPACE_DIRECTORY"] 169 requirements_file_tree = os.path.join(workspace, requirements_file_relative) 170 # In most cases, requirements_file will be a symlink to the real file in the source tree. 171 # If symlinks are not enabled (e.g. on Windows), then requirements_file will be a copy, 172 # and we should copy the updated requirements back to the source tree. 173 if not os.path.samefile(resolved_requirements_file, requirements_file_tree): 174 atexit.register( 175 lambda: shutil.copy( 176 resolved_requirements_file, requirements_file_tree 177 ) 178 ) 179 cli() 180 requirements_file_relative_path = Path(requirements_file_relative) 181 content = requirements_file_relative_path.read_text() 182 content = content.replace(absolute_path_prefix, "") 183 requirements_file_relative_path.write_text(content) 184 else: 185 # cli will exit(0) on success 186 try: 187 print("Checking " + requirements_file) 188 cli() 189 print("cli() should exit", file=sys.stderr) 190 sys.exit(1) 191 except SystemExit as e: 192 if e.code == 2: 193 print( 194 "pip-compile exited with code 2. This means that pip-compile found " 195 "incompatible requirements or could not find a version that matches " 196 f"the install requirement in {requirements_in_relative}.", 197 file=sys.stderr, 198 ) 199 sys.exit(1) 200 elif e.code == 0: 201 golden = open(_locate(bazel_runfiles, requirements_file)).readlines() 202 out = open(requirements_out).readlines() 203 out = [line.replace(absolute_path_prefix, "") for line in out] 204 if golden != out: 205 import difflib 206 207 print("".join(difflib.unified_diff(golden, out)), file=sys.stderr) 208 print( 209 "Lock file out of date. Run '" 210 + update_command 211 + "' to update.", 212 file=sys.stderr, 213 ) 214 sys.exit(1) 215 sys.exit(0) 216 else: 217 print( 218 f"pip-compile unexpectedly exited with code {e.code}.", 219 file=sys.stderr, 220 ) 221 sys.exit(1) 222