1#!/usr/bin/env python3 2# Copyright 2023 The Bazel Authors. All rights reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15 16"""A script to manage internal pip dependencies.""" 17 18from __future__ import annotations 19 20import argparse 21import json 22import os 23import pathlib 24import re 25import sys 26import tempfile 27import textwrap 28from dataclasses import dataclass 29 30from pip._internal.cli.main import main as pip_main 31 32from tools.private.update_deps.args import path_from_runfiles 33from tools.private.update_deps.update_file import update_file 34 35 36@dataclass 37class Dep: 38 name: str 39 url: str 40 sha256: str 41 42 43def _dep_snippet(deps: list[Dep]) -> str: 44 lines = [] 45 for dep in deps: 46 lines.extend( 47 [ 48 "(\n", 49 f' "{dep.name}",\n', 50 f' "{dep.url}",\n', 51 f' "{dep.sha256}",\n', 52 "),\n", 53 ] 54 ) 55 56 return textwrap.indent("".join(lines), " " * 4) 57 58 59def _module_snippet(deps: list[Dep]) -> str: 60 lines = [] 61 for dep in deps: 62 lines.append(f'"{dep.name}",\n') 63 64 return textwrap.indent("".join(lines), " " * 4) 65 66 67def _generate_report(requirements_txt: pathlib.Path) -> dict: 68 with tempfile.NamedTemporaryFile() as tmp: 69 tmp_path = pathlib.Path(tmp.name) 70 sys.argv = [ 71 "pip", 72 "install", 73 "--dry-run", 74 "--ignore-installed", 75 "--report", 76 f"{tmp_path}", 77 "-r", 78 f"{requirements_txt}", 79 ] 80 pip_main() 81 with open(tmp_path) as f: 82 return json.load(f) 83 84 85def _get_deps(report: dict) -> list[Dep]: 86 deps = [] 87 for dep in report["install"]: 88 try: 89 dep = Dep( 90 name="pypi__" 91 + re.sub( 92 "[._-]+", 93 "_", 94 dep["metadata"]["name"], 95 ), 96 url=dep["download_info"]["url"], 97 sha256=dep["download_info"]["archive_info"]["hash"][len("sha256=") :], 98 ) 99 except: 100 debug_dep = textwrap.indent(json.dumps(dep, indent=4), " " * 4) 101 print(f"Could not parse the response from 'pip':\n{debug_dep}") 102 raise 103 104 deps.append(dep) 105 106 return sorted(deps, key=lambda dep: dep.name) 107 108 109def main(): 110 parser = argparse.ArgumentParser(__doc__) 111 parser.add_argument( 112 "--start", 113 type=str, 114 default="# START: maintained by 'bazel run //tools/private/update_deps:update_pip_deps'", 115 help="The text to match in a file when updating them.", 116 ) 117 parser.add_argument( 118 "--end", 119 type=str, 120 default="# END: maintained by 'bazel run //tools/private/update_deps:update_pip_deps'", 121 help="The text to match in a file when updating them.", 122 ) 123 parser.add_argument( 124 "--dry-run", 125 action="store_true", 126 help="Wether to write to files", 127 ) 128 parser.add_argument( 129 "--requirements-txt", 130 type=path_from_runfiles, 131 default=os.environ.get("REQUIREMENTS_TXT"), 132 help="The requirements.txt path for the pypi tools, defaults to the value taken from REQUIREMENTS_TXT", 133 ) 134 parser.add_argument( 135 "--deps-bzl", 136 type=path_from_runfiles, 137 default=os.environ.get("DEPS_BZL"), 138 help="The path for the file to be updated, defaults to the value taken from DEPS_BZL", 139 ) 140 args = parser.parse_args() 141 142 report = _generate_report(args.requirements_txt) 143 deps = _get_deps(report) 144 145 update_file( 146 path=args.deps_bzl, 147 snippet=_dep_snippet(deps), 148 start_marker=args.start, 149 end_marker=args.end, 150 dry_run=args.dry_run, 151 ) 152 153 154if __name__ == "__main__": 155 main() 156