FIX: remove multiprocessing for gbk file filtering

Author: davidmeijer

scripts/filter_gbk_by_product.py

@@ -1,55 +1,65 @@
-"""Filter GenBank files by antiSMASH cand_cluster/protocluster product and copy matches."""
+#!/usr/bin/env python3
+"""Filter GenBank files by antiSMASH cand_cluster/protocluster product, EXACT product match, and copy matches."""
 
 from __future__ import annotations
 
 import argparse
 import os
 import re
 import shutil
-import sys
-from concurrent.futures import ProcessPoolExecutor
 from pathlib import Path
 from typing import Iterable, Sequence
 
 from tqdm import tqdm
 
 
+# Regex to capture all product qualifiers
+PRODUCT_RE = re.compile(r'/product="([^"]*)"', re.IGNORECASE)
+
+
 def _iter_gbk_files(root: Path, exts: Sequence[str]) -> Iterable[Path]:
-    for ext in exts:
-        yield from root.rglob(f"*{ext}")
+    """Yield all files with specified extensions."""
+    exts = {e.lower() for e in exts}
+    for p in root.rglob("*"):
+        if p.is_file() and p.suffix.lower() in exts:
+            yield p
 
 
 def _file_matches(
     gbk_path: str,
     products: Sequence[str],
 ) -> dict | None:
     """
-    Return match info if the GBK contains cand_cluster/protocluster with desired product.
-    Uses lightweight text search to avoid full parsing overhead.
+    Return match info if the GBK contains cand_cluster/protocluster and has a product EXACTLY equal
+    to one of the given search terms (case-insensitive).
     """
     path = Path(gbk_path)
+
     try:
         text = path.read_text(errors="ignore")
-    except Exception as exc:  # noqa: BLE001
+    except Exception as exc:
         return {"error": f"read failed: {exc}", "file": gbk_path}
 
     lower = text.lower()
+
+    # quick skip
     if "cand_cluster" not in lower and "protocluster" not in lower:
         return None
 
-    matched: set[str] = set()
-    for prod in products:
-        pat = re.compile(r'/product="[^"]*' + re.escape(prod.lower()) + r'[^"]*"', re.IGNORECASE)
-        if pat.search(text):
-            matched.add(prod)
+    # extract all /product="..."`
+    found_products = PRODUCT_RE.findall(text)
+    found_products_norm = {fp.strip().lower() for fp in found_products}
+
+    # find exact matches
+    matched = sorted(fp for fp in found_products_norm if fp in products)
 
     if not matched:
         return None
 
     return {
         "file": path.name,
         "path": str(path),
-        "matched_products": sorted(matched),
+        "matched_products": matched,
     }
 
 
@@ -62,34 +72,23 @@ def cli() -> argparse.Namespace:
         "--products",
         nargs="+",
         required=True,
-        help='Product substrings to match (e.g., "NRPS", "T1PKS")',
-    )
-    parser.add_argument(
-        "-w",
-        "--workers",
-        type=int,
-        default=os.cpu_count() or 1,
-        help="Number of worker processes",
+        help='Exact product terms to match (case-insensitive)',
     )
     parser.add_argument(
         "--exts",
         nargs="+",
         default=[".gbk", ".gbff", ".gb"],
         help="File extensions to include",
     )
-    parser.add_argument(
-        "--chunksize",
-        type=int,
-        default=16,
-        help="Chunksize for process pool mapping",
-    )
     return parser.parse_args()
 
 
 def main() -> None:
     args = cli()
     input_dir: Path = args.input_dir
     output_dir: Path = args.output_dir
+
+    # normalize to lowercase for exact matching
     products = [p.lower() for p in args.products]
 
     if not input_dir.exists():
@@ -101,38 +100,31 @@ def main() -> None:
 
     output_dir.mkdir(parents=True, exist_ok=True)
 
-    total = len(files)
     copied = 0
     errors = 0
 
-    with ProcessPoolExecutor(max_workers=args.workers) as executor:
-        with tqdm(total=total, desc="Scanning GBKs", unit="file") as pbar:
-            for result in executor.map(
-                _file_matches,
-                (str(p) for p in files),
-                [products] * len(files),
-                chunksize=args.chunksize,
-            ):
-                pbar.update(1)
-
-                if not result:
-                    continue
-
-                if isinstance(result, dict) and "error" in result:
-                    errors += 1
-                    pbar.set_postfix(copied=copied, errors=errors)
-                    continue
-
-                src_path = Path(result["path"])
-                dst_path = output_dir / src_path.name
-                try:
-                    shutil.copy2(src_path, dst_path)
-                    copied += 1
-                except Exception as exc:  # noqa: BLE001
-                    errors += 1
-                    result = {"error": f"copy failed: {exc}", "file": str(src_path)}
+    with tqdm(total=len(files), desc="Scanning GBKs", unit="file") as pbar:
+        for path in files:
+            result = _file_matches(str(path), products)
+            pbar.update(1)
+
+            if not result:
+                continue
 
+            if "error" in result:
+                errors += 1
                 pbar.set_postfix(copied=copied, errors=errors)
+                continue
+
+            try:
+                shutil.copy2(result["path"], output_dir / result["file"])
+                copied += 1
+            except Exception as exc:
+                errors += 1
+
+            pbar.set_postfix(copied=copied, errors=errors)
+
+    print(f"\nDone. Copied: {copied}, errors: {errors}, total scanned: {len(files)}")
 
 
 if __name__ == "__main__":