[llvm] [utils] Add script to generate elaborated assembly tests (PR #89026)

[Fangrui Song via llvm-commits]

================
@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Generate test body using split-file and a custom script.
+Currently, only assembly files are supported by placing generation instructions
+surrounded by .ifdef GEN/.endif directives.
+
+    .ifdef GEN
+    #--- a.cc
+    int va;
+    #--- gen
+    clang --target=aarch64-linux -S -g a.cc -o -
+    .endif
+    # content generated by the script 'gen'
+
+The script will prepare extra files with `split-file`, invoke `gen`, and then
+rewrite the part after `.endif` with its stdout.
+
+Example:
+PATH=/path/to/clang_build/bin:$PATH llvm/utils/update_test_body.py path/to/test.s
+"""
+import argparse, contextlib, os, subprocess, sys, tempfile
+
+
+ at contextlib.contextmanager
+def cd(dir):
+    cwd = os.getcwd()
+    os.chdir(dir)
+    try:
+        yield
+    finally:
+        os.chdir(cwd)
+
+
+def process(args, path):
+    split_file_input = []
+    prolog = []
+    is_split_file_input = False
+    is_prolog = True
+    with open(path) as f:
+        for line in f.readlines():
+            line = line.rstrip()
+            if is_prolog:
+                prolog.append(line)
+            if line.startswith(".endif"):
+                is_split_file_input = is_prolog = False
+            if is_split_file_input:
+                split_file_input.append(line)
+            if line.startswith(".ifdef GEN"):
+                is_split_file_input = True
+
+    if not split_file_input:
+        print("no .ifdef GEN", file=sys.stderr)
+        return
+    if is_split_file_input:
+        print("no .endif", file=sys.stderr)
+        return
+    with tempfile.TemporaryDirectory(prefix="update_test_body_") as dir:
+        sub = subprocess.run(
+            ["split-file", "-", dir],
+            input="\n".join(split_file_input).encode(),
+            capture_output=True,
+        )
+        if sub.returncode != 0:
+            sys.stderr.write(f"split-file failed\n{sub.stderr.decode()}")
+            return
+        with cd(dir):
+            if args.shell:
+                print(f"invoke shell in the temporary directory '{dir}'")
+                subprocess.run([os.environ.get("SHELL", "sh")])
+                return
+            if not os.path.exists("gen"):
+                print("'gen' does not exist", file=sys.stderr)
+                return
+
+            sub = subprocess.run(
+                ["sh", "-euo", "pipefail", "gen"],
+                capture_output=True,
+                # Don't encode the directory information to the Clang output.
+                # Remove unneeded details (.ident) as well.
+                env=dict(
+                    os.environ,
+                    CCC_OVERRIDE_OPTIONS="#+-fno-ident",
+                    PWD="/proc/self/cwd",
+                ),
+            )
+            sys.stderr.write(sub.stderr.decode())
+            if sub.returncode != 0:
+                print("'gen' failed", file=sys.stderr)
----------------
MaskRay wrote:

`sys.stderr.write` does not append a newline while `print` does by default (unless `end=''`)... I feel that sys.stderr.write is quite common as well.

https://github.com/llvm/llvm-project/pull/89026