refactor: clean up codeql_test_run.py

Author: redsun82

misc/just/codeql_test_run.py

@@ -15,6 +15,7 @@
 import re
 import subprocess
 import sys
+import shutil
 from pathlib import Path
 
 JUST = os.environ.get("JUST_EXECUTABLE", "just")
@@ -23,12 +24,12 @@
 CMD_END = os.environ.get("CMD_END", "")
 SEMMLE_CODE = os.environ.get("SEMMLE_CODE")
 
-ENV_RE = re.compile(r"^[A-Z_][A-Z_0-9]*=.*$")
+ENV_RE = re.compile(r"(^[A-Z_][A-Z_0-9]*)=(.*)$")
 
 
 def invoke(invocation, *, cwd=None, log_prefix=""):
     prefix = f"{log_prefix} " if log_prefix else ""
-    print(f"{CMD_BEGIN}{prefix}{' '.join(invocation)}{CMD_END}")
+    print(f"{CMD_BEGIN}{prefix}{' '.join(map(str, invocation))}{CMD_END}")
     try:
         subprocess.run(invocation, check=True, cwd=cwd)
     except subprocess.CalledProcessError as e:
@@ -61,11 +62,11 @@ def build_parser():
     parser.add_argument("--codeql", default="build" if SEMMLE_CODE else "host")
     parser.add_argument("--extra-check", action="append", dest="extra_checks")
     parser.add_argument("--all-checks", "+", action="store_true", dest="all")
-    parser.set_defaults(extra_checks=[], tests=[], flags=[], env=[])
+    parser.set_defaults(extra_checks=[], tests=[], flags=[], env={})
     return parser
 
 
-def parse_arguments(argv):
+def parse_arguments():
     """Sort a command line into the kinds that are handled differently.
 
     `argparse` owns the options this script acts on itself. Everything else belongs to
@@ -74,33 +75,44 @@ def parse_arguments(argv):
     positionals, told apart only by how they look.
     """
     # An empty argument can come from a caller interpolating an unset variable.
-    args, rest = build_parser().parse_known_args([arg for arg in argv if arg])
+    args, rest = build_parser().parse_known_args()
     for arg in rest:
-        if arg.startswith("-"):
+        if not arg:
+            pass
+        elif arg.startswith("-"):
             args.flags.append(arg)
-        elif ENV_RE.match(arg):
-            args.env.append(arg)
+        elif m := ENV_RE.match(arg):
+            k, v = m.groups()
+            args.env[k] = v
         else:
             args.tests.append(arg)
     return args
 
 
-def env_value(args, name, default):
-    """Resolve a setting from test arguments, then the environment, then a default."""
-    for assignment in reversed(args.env):
-        key, _, value = assignment.partition("=")
-        if key == name and value:
-            return value
-    return os.environ.get(name) or default
-
+def resolve_codeql(args: argparse.Namespace) -> Path:
+    suffix = ".exe" if sys.platform == "win32" else ""
+    match args.codeql:
+        case "built" | "build":
+            return Path(
+                SEMMLE_CODE, "target", "intree", f"codeql-{args.language}", "codeql"
+            ).with_suffix(suffix)
+        case "host":
+            return Path(shutil.which("codeql" + suffix))
+        case _:
+            codeql = Path(args.codeql)
+            if codeql.is_dir():
+                codeql /= "codeql"
+            return codeql.with_suffix(suffix)
+        
 
 def main():
-    argv = sys.argv[1:]
-    args = parse_arguments(argv)
+    args = parse_arguments()
+
     if args.all:
         # Apply what the language offered by parsing it alongside everything else, so an
         # offered check lands exactly where the same flag typed by hand would.
-        args = parse_arguments(argv + args.extra_checks)
+        sys.argv[1:1] = args.extra_checks
+        args = parse_arguments()
 
     if not SEMMLE_CODE and args.codeql in ("build", "built"):
         error(
@@ -112,52 +124,33 @@ def main():
     if not args.tests:
         args.tests.append(".")
 
+    os.environ.update(args.env)
+
     # Resolve these only once all arguments are known, so that a `RAM_PER_THREAD=` test
     # argument can lower the default on memory-heavy suites.
     default_ram = 3000 if sys.platform == "linux" else 2048
-    ram_per_thread = int(env_value(args, "RAM_PER_THREAD", default_ram))
-    cpus = int(env_value(args, "CPUS", os.cpu_count() or 1))
+    ram_per_thread = int(os.environ.get("RAM_PER_THREAD", default_ram))
+    cpus = int(os.environ.get("CPUS", os.cpu_count() or 1))
     args.flags[:0] = [f"--ram={ram_per_thread * cpus}", f"-j{cpus}"]
 
     if args.codeql == "build":
-        if invoke([JUST, args.language, "build"], cwd=SEMMLE_CODE) != 0:
-            return 1
+        if ret := invoke([JUST, args.language, "build"], cwd=SEMMLE_CODE):
+            return ret
 
     if args.codeql != "host":
         # Disable the default implicit config file, but keep an explicit one.
         # Same behavior wrt --codeql as the integration test runner.
         os.environ.setdefault("CODEQL_CONFIG_FILE", ".")
 
-    for env_var in args.env:
-        key, _, value = env_var.partition("=")
-        os.environ[key] = value
-
-    # Resolve codeql executable
-    if args.codeql in ("built", "build"):
-        codeql = Path(
-            SEMMLE_CODE, "target", "intree", f"codeql-{args.language}", "codeql"
-        )
-    elif args.codeql == "host":
-        codeql = Path("codeql")
-    else:
-        codeql = Path(args.codeql)
-
-    if codeql.is_dir():
-        codeql = codeql / "codeql"
-
-    # On Windows, prefer codeql.exe over the Unix shell wrapper
-    if sys.platform == "win32" and codeql.suffix != ".exe":
-        exe = codeql.with_suffix(".exe")
-        if exe.exists():
-            codeql = exe
+    codeql = resolve_codeql(args)
 
-    if args.codeql != "host" and not codeql.exists():
+    if not codeql.exists():
         error(f"CodeQL executable not found: {codeql}")
         return 1
 
     return invoke(
-        [str(codeql), "test", "run", *args.flags, "--", *args.tests],
-        log_prefix=" ".join(args.env),
+        [codeql, "test", "run", *args.flags, "--", *args.tests],
+        log_prefix=" ".join(f"{k}={v}" for k, v in args.env.items()),
     )