llvm-project

1
#!/usr/bin/env python3
2
"""Generate test body using split-file and a custom script.
3

4
The script will prepare extra files with `split-file`, invoke `gen`, and then
5
rewrite the part after `gen` with its stdout.
6

7
https://llvm.org/docs/TestingGuide.html#elaborated-tests
8

9
Example:
10
PATH=/path/to/clang_build/bin:$PATH llvm/utils/update_test_body.py path/to/test.s
11
"""
12
import argparse
13
import contextlib
14
import os
15
import re
16
import subprocess
17
import sys
18
import tempfile
19

20

21
@contextlib.contextmanager
22
def cd(directory):
23
    cwd = os.getcwd()
24
    os.chdir(directory)
25
    try:
26
        yield
27
    finally:
28
        os.chdir(cwd)
29

30

31
def process(args, path):
32
    prolog = []
33
    seen_gen = False
34
    with open(path) as f:
35
        for line in f.readlines():
36
            line = line.rstrip()
37
            prolog.append(line)
38
            if (seen_gen and re.match(r"(.|//)---", line)) or line.startswith(".endif"):
39
                break
40
            if re.match(r"(.|//)--- gen", line):
41
                seen_gen = True
42
        else:
43
            print(
44
                "'gen' should be followed by another part (---) or .endif",
45
                file=sys.stderr,
46
            )
47
            return 1
48

49
    if not seen_gen:
50
        print("'gen' does not exist", file=sys.stderr)
51
        return 1
52
    with tempfile.TemporaryDirectory(prefix="update_test_body_") as dir:
53
        try:
54
            # If the last line starts with ".endif", remove it.
55
            sub = subprocess.run(
56
                ["split-file", "-", dir],
57
                input="\n".join(
58
                    prolog[:-1] if prolog[-1].startswith(".endif") else prolog
59
                ).encode(),
60
                capture_output=True,
61
                check=True,
62
            )
63
        except subprocess.CalledProcessError as ex:
64
            sys.stderr.write(ex.stderr.decode())
65
            return 1
66
        with cd(dir):
67
            if args.shell:
68
                print(f"invoke shell in the temporary directory '{dir}'")
69
                subprocess.run([os.environ.get("SHELL", "sh")])
70
                return 0
71

72
            sub = subprocess.run(
73
                ["sh", "-eu", "gen"],
74
                capture_output=True,
75
                # Don't encode the directory information to the Clang output.
76
                # Remove unneeded details (.ident) as well.
77
                env=dict(
78
                    os.environ,
79
                    CCC_OVERRIDE_OPTIONS="#^-fno-ident",
80
                    PWD="/proc/self/cwd",
81
                ),
82
            )
83
            sys.stderr.write(sub.stderr.decode())
84
            if sub.returncode != 0:
85
                print("'gen' failed", file=sys.stderr)
86
                return sub.returncode
87
            if not sub.stdout:
88
                print("stdout is empty; forgot -o - ?", file=sys.stderr)
89
                return 1
90
            content = sub.stdout.decode()
91

92
    with open(path, "w") as f:
93
        # Print lines up to '.endif'.
94
        print("\n".join(prolog), file=f)
95
        # Then print the stdout of 'gen'.
96
        f.write(content)
97

98

99
parser = argparse.ArgumentParser(
100
    description="Generate test body using split-file and a custom script"
101
)
102
parser.add_argument("files", nargs="+")
103
parser.add_argument(
104
    "--shell", action="store_true", help="invoke shell instead of 'gen'"
105
)
106
args = parser.parse_args()
107
for path in args.files:
108
    retcode = process(args, path)
109
    if retcode != 0:
110
        sys.exit(retcode)
111