llvm/utils/update_test_body.py

   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)