fix(clew): log completed AST extractions

src/clew/clew.py

@@ -8,11 +8,13 @@ import subprocess
 import sys
 from collections import deque
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
 from rpycdec.rpa import extract_rpa
-from src.clew import extract_renpy_ast as dump_ast_script
+from src.clew import extract_renpy_ast as extract_renpy_ast_script
 
+EXTRACTOR = Path(extract_renpy_ast_script.__file__).resolve()
 LOG = logging.getLogger(__name__)
 
 def parse_args(argv: list[str]) -> argparse.Namespace:
@@ -97,16 +99,19 @@ def find(root: Path, predicate: Callable[[Path], bool]) -> list[Path]:
                 matches.append(file_path)
     return sorted(matches)
 
+def is_game_root(path: Path) -> bool:
+    return (
+        (path / "game").is_dir()
+        and (path / "renpy").is_dir()
+        and (path / "lib" / "linux-x86_64" / "python").is_file())
 
-def ast_dump_output_path(output_root: Path, extracted_root: Path, game_root: Path, script: Path) -> Path:
-    namespace = game_root.relative_to(extracted_root)
-    relative_script = script.relative_to(game_root)
-    return output_root / "ast" / namespace / relative_script.with_suffix(relative_script.suffix + ".json")
+
+def is_renpy_script(path: Path) -> bool:
+    return path.is_file() and path.suffix.lower() in {".rpyc", ".rpymc"}
 
 
 def run_extract_renpy_ast(
     bundled_python: Path,
-    extractor: Path,
     game_root: Path,
     compiled_script: Path,
     output: Path,
@@ -115,7 +120,7 @@ def run_extract_renpy_ast(
     command = [
         "steam-run",
         str(bundled_python),
-        str(extractor),
+        str(EXTRACTOR),
         "--game-root",
         str(game_root),
         "--input",
@@ -126,34 +131,42 @@ def run_extract_renpy_ast(
     subprocess.run(command, check=True)
 
 
-def extract_game_ast(output_root: Path, extracted_root: Path, overwrite: bool) -> None:
-    extractor = Path(dump_ast_script.__file__).resolve()
-    game_roots = find(
-        extracted_root,
-        lambda path: (
-            (path / "game").is_dir()
-            and (path / "renpy").is_dir()
-            and (path / "lib" / "linux-x86_64" / "python").is_file()
-        ),
-    )
+def extract_game_ast(extracted_root: Path, ast_root: Path, overwrite: bool) -> None:
+    game_roots = find(extracted_root, is_game_root)
     if not game_roots:
-        LOG.warning("no Ren'Py game roots found under %s", extracted_root)
+        LOG.warning(f"no Ren'Py game roots found under {extracted_root}")
         return
 
-    for game_root in game_roots:
-        bundled_python = game_root / "lib" / "linux-x86_64" / "python"
-        scripts = find(
-            game_root / "game",
-            lambda path: path.is_file() and path.suffix.lower() in {".rpyc", ".rpymc"},
-        )
-        LOG.info("found %d compiled Ren'Py scripts under %s", len(scripts), game_root)
-        for script in scripts:
-            output = ast_dump_output_path(output_root, extracted_root, game_root, script)
-            if output.exists() and not overwrite:
-                LOG.info("skipping existing AST dump: %s", output)
+    tasks = [
+        (game_root, script_path)
+        for game_root in game_roots
+        for script_path in find(game_root / "game", is_renpy_script)
+    ]
+
+    with ThreadPoolExecutor() as executor:
+        futures = {}
+        for game_root, script_path in tasks:
+            output_path = (
+                ast_root
+                / game_root.relative_to(extracted_root)
+                / script_path.relative_to(game_root).with_suffix(script_path.suffix + ".json")
+            )
+            if output_path.exists() and not overwrite:
+                LOG.info(f"skipping existing AST extraction: {output_path}")
                 continue
-            LOG.info("dumping AST: %s", script)
-            run_extract_renpy_ast(bundled_python, extractor, game_root, script, output)
+
+            future = executor.submit(
+                run_extract_renpy_ast,
+                game_root / "lib" / "linux-x86_64" / "python",
+                game_root,
+                script_path,
+                output_path,
+            )
+            futures[future] = script_path
+
+        for future in as_completed(futures):
+            future.result()
+            LOG.info(f"finished AST extraction: {futures[future]}")
 
 
 def run(args: argparse.Namespace) -> None:
@@ -168,7 +181,8 @@ def run(args: argparse.Namespace) -> None:
     extract_game_files(game_root, extracted)
 
     if not args.no_extract_renpy_ast:
-        extract_game_ast(output_root, extracted, args.overwrite)
+        ast = get_output_directory(output_root, "ast")
+        extract_game_ast(extracted, ast, args.overwrite)
 
 
 if __name__ == "__main__":