wandb

1
#!/usr/bin/env python
2
"""Generate code for wandb SDK.
3

4
Usage:
5
    ./tools/generate-tool.py --generate
6
    ./tools/generate-tool.py --check
7
    ./tools/generate-tool.py --generate --check
8
    ./tools/generate-tool.py --generate --check wandb/sdk/lib/_wburls_generated.py
9
"""
10

11
import argparse
12
import contextlib
13
import filecmp
14
import os
15
import subprocess
16
import tempfile
17
from pathlib import Path, PurePath
18
from typing import Iterator, List
19

20
parser = argparse.ArgumentParser()
21
parser.add_argument("--generate", action="store_true", help="generate files")
22
parser.add_argument(
23
    "--no-format",
24
    dest="format",
25
    action="store_false",
26
    default=True,
27
    help="format files",
28
)
29
parser.add_argument(
30
    "--check", action="store_true", help="check if generated files are up-to-date"
31
)
32
parser.add_argument("files", nargs="*")
33
args = parser.parse_args()
34

35

36
GENERATE_SUFFIX = "_generate.py"
37
GENERATED_SUFFIX = "_generated.py"
38

39

40
def get_paths() -> List[PurePath]:
41
    paths: List[PurePath] = []
42
    if not args.files:
43
        exclude_dirs = {"vendor", "__pycache__"}
44
        root_dir = Path(__file__).resolve().parent.parent / "wandb"
45
        for base, subdirs, files in os.walk(root_dir):
46
            # Don't walk into excluded subdirectories
47
            subdirs[:] = list(set(subdirs) - exclude_dirs)
48
            for fname in files:
49
                if fname.endswith(GENERATE_SUFFIX):
50
                    paths.append(PurePath(base, fname))
51
    for f in args.files:
52
        paths.append(Path(f))
53
    return paths
54

55

56
def generate_file(generate_path: Path, output_path: Path) -> None:
57
    status, output = subprocess.getstatusoutput(f"python {generate_path}")
58
    assert status == 0, f"Error: {output}"
59
    with open(output_path, "w") as f:
60
        f.write("# DO NOT EDIT -- GENERATED BY: `generate-tool.py --generate`")
61
        f.write(output)
62

63

64
def generate_files(paths: List[PurePath]) -> None:
65
    for p in paths:
66
        output_path = p.parent / str(p).replace(GENERATE_SUFFIX, GENERATED_SUFFIX)
67
        print(f"INFO: Generating {output_path}...")
68
        generate_file(p, output_path)
69

70

71
def format_file(filename: Path) -> None:
72
    status, output = subprocess.getstatusoutput(f"black {filename}")
73
    assert status == 0, f"Error: {output}"
74

75

76
def format_files(paths: List[PurePath]) -> None:
77
    for p in paths:
78
        output_path = p.parent / str(p).replace(GENERATE_SUFFIX, GENERATED_SUFFIX)
79
        print(f"INFO: Formatting {output_path}...")
80
        format_file(output_path)
81

82

83
@contextlib.contextmanager
84
def temp_fname() -> Iterator[PurePath]:
85
    try:
86
        f = tempfile.NamedTemporaryFile(delete=False)
87
        tmp_name = f.name
88
        f.close()
89
        yield Path(tmp_name)
90
    finally:
91
        os.unlink(tmp_name)
92

93

94
def check_files(paths: List[PurePath]) -> None:
95
    for p in paths:
96
        generated_path = p.parent / str(p).replace(GENERATE_SUFFIX, GENERATED_SUFFIX)
97
        print(f"INFO: Checking {generated_path}...")
98
        with temp_fname() as temp_file:
99
            generate_file(p, temp_file)
100
            format_file(temp_file)
101
            assert filecmp.cmp(
102
                generated_path, temp_file
103
            ), f"expected: {open(temp_file).read()}"
104

105

106
def main() -> None:
107
    path_list = get_paths()
108

109
    if args.generate:
110
        generate_files(path_list)
111
        if args.format:
112
            format_files(path_list)
113

114
    if args.check:
115
        check_files(path_list)
116

117

118
if __name__ == "__main__":
119
    main()
120