feat(clew): extract Ren'Py ASTs

src/clew/clew.py

@@ -4,12 +4,17 @@ import argparse
 import logging
 import os
 import shutil
+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 extract_renpy_ast_script
 
+EXTRACTOR = Path(extract_renpy_ast_script.__file__).resolve()
 LOG = logging.getLogger(__name__)
 
 def parse_args(argv: list[str]) -> argparse.Namespace:
@@ -34,6 +39,11 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
         action="store_true",
         help="Replace Clew-managed output directories.",
     )
+    parser.add_argument(
+        "--no-extract-renpy-ast",
+        action="store_true",
+        help="Skip Ren'Py AST extraction.",
+    )
     return parser.parse_args(argv)
 
 
@@ -60,7 +70,7 @@ def extract_archive(archive: Path, destination: Path) -> None:
         extract_rpa(handle, output_dir=str(destination))
 
 
-def extract_game_files(game_root: Path, output_path: Path) -> None:
+def extract_game_files(game_root: Path, output_path: Path, overwrite: bool) -> None:
     queue = deque(game_root.iterdir())
 
     while queue:
@@ -71,22 +81,112 @@ def extract_game_files(game_root: Path, output_path: Path) -> None:
             destination.mkdir(parents=True, exist_ok=True)
             queue.extend(item.iterdir())
         elif item.is_file() and item.suffix.lower() == ".rpa":
-            extract_archive(item, destination.with_suffix(""))
+            archive_output = destination.with_suffix("")
+            if archive_output.is_dir() and not overwrite:
+                LOG.info(f"skipping existing archive extraction: {archive_output}")
+                continue
+            extract_archive(item, archive_output)
         elif item.is_file():
             destination.parent.mkdir(parents=True, exist_ok=True)
             shutil.copy2(item, destination)
 
 
+def find(root: Path, predicate: Callable[[Path], bool]) -> list[Path]:
+    matches = []
+    for current, _directories, files in os.walk(root):
+        current_path = Path(current)
+        if predicate(current_path):
+            matches.append(current_path)
+        for file_name in files:
+            file_path = current_path / file_name
+            if predicate(file_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 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,
+    game_root: Path,
+    compiled_script: Path,
+    output: Path,
+) -> None:
+    output.parent.mkdir(parents=True, exist_ok=True)
+    command = [
+        "steam-run",
+        str(bundled_python),
+        str(EXTRACTOR),
+        "--game-root",
+        str(game_root),
+        "--input",
+        str(compiled_script),
+        "--output",
+        str(output),
+    ]
+    subprocess.run(command, check=True)
+
+
+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(f"no Ren'Py game roots found under {extracted_root}")
+        return
+
+    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
+
+            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:
     game_root = args.game_root.resolve()
     output_root = resolve_output_root(args.output_root)
 
     if not game_root.is_dir():
-        LOG.critical(f'game root is not a directory: {game_root}')
+        LOG.critical(f"game root is not a directory: {game_root}")
         raise SystemExit(2)
 
     extracted = get_output_directory(output_root, "extracted")
-    extract_game_files(game_root, extracted)
+    extract_game_files(game_root, extracted, args.overwrite)
+
+    if not args.no_extract_renpy_ast:
+        ast = get_output_directory(output_root, "ast")
+        extract_game_ast(extracted, ast, args.overwrite)
 
 
 if __name__ == "__main__":