pytorch

1
# mypy: allow-untyped-defs
2
import argparse
3
import glob
4
import os
5
from pathlib import Path
6
from 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.
11
DENY_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

29
strip_file_dir = ""
30

31

32
def remove_prefix(text, prefix):
33
    if text.startswith(prefix):
34
        return text[len(prefix) :]
35
    return text
36

37

38
def 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

46
def 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

85
if __name__ == "__main__":
86
    main()  # pragma: no cover
87