xref: /llvm-project/llvm/utils/update_test_body.py (revision aacea0d0f67401f5a0b74947f3ff179ade9cbf6d) 
1#!/usr/bin/env python3
2"""Generate test body using split-file and a custom script.
3
4The script will prepare extra files with `split-file`, invoke `gen`, and then
5rewrite the part after `gen` with its stdout.
6
7https://llvm.org/docs/TestingGuide.html#elaborated-tests
8
9Example:
10PATH=/path/to/clang_build/bin:$PATH llvm/utils/update_test_body.py path/to/test.s
11"""
12import argparse
13import contextlib
14import os
15import re
16import subprocess
17import sys
18import tempfile
19
20
21@contextlib.contextmanager
22def cd(directory):
23    cwd = os.getcwd()
24    os.chdir(directory)
25    try:
26        yield
27    finally:
28        os.chdir(cwd)
29
30
31def 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
99parser = argparse.ArgumentParser(
100    description="Generate test body using split-file and a custom script"
101)
102parser.add_argument("files", nargs="+")
103parser.add_argument(
104    "--shell", action="store_true", help="invoke shell instead of 'gen'"
105)
106args = parser.parse_args()
107for path in args.files:
108    retcode = process(args, path)
109    if retcode != 0:
110        sys.exit(retcode)
111