Analyzer cleanup

tools/analyze.py

@@ -1,12 +1,12 @@
 import math
 import os
 import time
+import json
 from collections import defaultdict
-from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import ProcessPoolExecutor, Executor
 from pathlib import Path
 from dataclasses import dataclass, asdict
 from itertools import combinations
-from typing import Callable
 
 import numpy as np
 from PIL import Image
@@ -192,29 +192,30 @@ def output_db(results: list[dict], db_file: str):
     output.write("]\n")
 
 
-if __name__ == "__main__":
-  from sys import argv
-  dex_file = argv[1] if len(argv) > 1 else "data/pokedex.json"
-  image_dir = argv[2] if len(argv) > 2 else "images"
-  seed = int(argv[3]) if len(argv) > 3 else 230308
-  db_file = argv[4] if len(argv) > 4 else "data/latest.db"
-
-  import json
-  with open(dex_file) as infile:
-    dex = json.load(infile)
-
+def run_ingest(dex: dict, image_dir: str, exec: Executor, db_file: str):
   ingest = Ingester(dex)
   to_process = search_files(image_dir)
-
   start = time.time()
-  with ProcessPoolExecutor(4) as pool:
-    results = list(pool.map(ingest, to_process.items()))
+  results = list(exec.map(ingest, to_process.items()))
   end = time.time()
   print(
     f"Finished ingest of {len(to_process)} forms",
     f"and {sum(len(fns) for fns in to_process.values())} files",
     f"in {(end - start):.2f}s"
   )
-
   output_db(sorted(results, key=lambda e: (e["num"], e["name"])), db_file)
   print(f"Output {len(results)} entries to {db_file}")
+
+
+if __name__ == "__main__":
+  from sys import argv
+  dex_file = argv[1] if len(argv) > 1 else "data/pokedex.json"
+  image_dir = argv[2] if len(argv) > 2 else "images"
+  seed = int(argv[3]) if len(argv) > 3 else 230308
+  db_file = argv[4] if len(argv) > 4 else "data/latest.db"
+
+  with open(dex_file) as infile:
+    dex = json.load(infile)
+
+  with ProcessPoolExecutor(4) as pool:
+    run_ingest(dex, image_dir, pool, db_file)