xref: /aosp_15_r20/external/pytorch/torch/utils/_zip.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# mypy: allow-untyped-defs
2import argparse
3import glob
4import os
5from pathlib import Path
6from zipfile import ZipFile
7
8# Exclude some standard library modules to:
9# 1. Slim down the final zipped file size
10# 2. Remove functionality we don't want to support.
11DENY_LIST = [
12    # Interface to unix databases
13    "dbm",
14    # ncurses bindings (terminal interfaces)
15    "curses",
16    # Tcl/Tk GUI
17    "tkinter",
18    "tkinter",
19    # Tests for the standard library
20    "test",
21    "tests",
22    "idle_test",
23    "__phello__.foo.py",
24    # importlib frozen modules. These are already baked into CPython.
25    "_bootstrap.py",
26    "_bootstrap_external.py",
27]
28
29strip_file_dir = ""
30
31
32def remove_prefix(text, prefix):
33    if text.startswith(prefix):
34        return text[len(prefix) :]
35    return text
36
37
38def write_to_zip(file_path, strip_file_path, zf, prepend_str=""):
39    stripped_file_path = prepend_str + remove_prefix(file_path, strip_file_dir + "/")
40    path = Path(stripped_file_path)
41    if path.name in DENY_LIST:
42        return
43    zf.write(file_path, stripped_file_path)
44
45
46def main() -> None:
47    global strip_file_dir
48    parser = argparse.ArgumentParser(description="Zip py source")
49    parser.add_argument("paths", nargs="*", help="Paths to zip.")
50    parser.add_argument(
51        "--install-dir", "--install_dir", help="Root directory for all output files"
52    )
53    parser.add_argument(
54        "--strip-dir",
55        "--strip_dir",
56        help="The absolute directory we want to remove from zip",
57    )
58    parser.add_argument(
59        "--prepend-str",
60        "--prepend_str",
61        help="A string to prepend onto all paths of a file in the zip",
62        default="",
63    )
64    parser.add_argument("--zip-name", "--zip_name", help="Output zip name")
65
66    args = parser.parse_args()
67
68    zip_file_name = args.install_dir + "/" + args.zip_name
69    strip_file_dir = args.strip_dir
70    prepend_str = args.prepend_str
71    zf = ZipFile(zip_file_name, mode="w")
72
73    for p in sorted(args.paths):
74        if os.path.isdir(p):
75            files = glob.glob(p + "/**/*.py", recursive=True)
76            for file_path in sorted(files):
77                # strip the absolute path
78                write_to_zip(
79                    file_path, strip_file_dir + "/", zf, prepend_str=prepend_str
80                )
81        else:
82            write_to_zip(p, strip_file_dir + "/", zf, prepend_str=prepend_str)
83
84
85if __name__ == "__main__":
86    main()  # pragma: no cover
87