xref: /aosp_15_r20/external/pytorch/tools/linter/adapters/pip_init.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1"""
2Initializer script that installs stuff to pip.
3"""
4
5from __future__ import annotations
6
7import argparse
8import logging
9import os
10import shutil
11import subprocess
12import sys
13import time
14
15
16def run_command(args: list[str]) -> subprocess.CompletedProcess[bytes]:
17    logging.debug("$ %s", " ".join(args))
18    start_time = time.monotonic()
19    try:
20        return subprocess.run(args, check=True)
21    finally:
22        end_time = time.monotonic()
23        logging.debug("took %dms", (end_time - start_time) * 1000)
24
25
26if __name__ == "__main__":
27    parser = argparse.ArgumentParser(description="pip initializer")
28    parser.add_argument(
29        "packages",
30        nargs="+",
31        help="pip packages to install",
32    )
33    parser.add_argument(
34        "--verbose",
35        action="store_true",
36        help="verbose logging",
37    )
38    parser.add_argument(
39        "--dry-run", help="do not install anything, just print what would be done."
40    )
41    parser.add_argument(
42        "--no-black-binary",
43        help="do not use pre-compiled binaries from pip for black.",
44        action="store_true",
45    )
46
47    args = parser.parse_args()
48
49    logging.basicConfig(
50        format="<%(threadName)s:%(levelname)s> %(message)s",
51        level=logging.NOTSET if args.verbose else logging.DEBUG,
52        stream=sys.stderr,
53    )
54
55    uv_available = shutil.which("uv") is not None
56
57    if uv_available:
58        pip_args = ["uv", "pip", "install"]
59    else:
60        pip_args = ["pip", "install"]
61
62    # If we are in a global install, use `--user` to install so that you do not
63    # need root access in order to initialize linters.
64    #
65    # However, `pip install --user` interacts poorly with virtualenvs (see:
66    # https://bit.ly/3vD4kvl) and conda (see: https://bit.ly/3KG7ZfU). So in
67    # these cases perform a regular installation.
68    in_conda = os.environ.get("CONDA_PREFIX") is not None
69    in_virtualenv = os.environ.get("VIRTUAL_ENV") is not None
70    if not in_conda and not in_virtualenv:
71        pip_args.append("--user")
72
73    pip_args.extend(args.packages)
74
75    for package in args.packages:
76        package_name, _, version = package.partition("=")
77        if version == "":
78            raise RuntimeError(
79                "Package {package_name} did not have a version specified. "
80                "Please specify a version to produce a consistent linting experience."
81            )
82        if args.no_black_binary and "black" in package_name:
83            pip_args.append(f"--no-binary={package_name}")
84
85    dry_run = args.dry_run == "1"
86    if dry_run:
87        print(f"Would have run: {pip_args}")
88        sys.exit(0)
89
90    run_command(pip_args)
91