feat: per-cluster import extraction + JSON output

- _extract_imports scans source files for language-agnostic import patterns
- analyze_clusters tracks file paths through frequency_filter via identity
- Each cluster output includes its unique imports and file paths
- --json flag outputs structured JSON for LLM prompt injection
- Flat mode (--cluster false) also extracts imports

bex/tag_preprocessor/analyze.py

@@ -7,7 +7,9 @@ Runs the full Phase 1.0 pipeline over a directory of source files.
 """
 
 import argparse
+import json
 import os
+import re
 import sys
 from pathlib import Path, PurePath
 from collections import Counter
@@ -17,6 +19,16 @@ import pathspec
 from .code import preprocess_by_method, _extract_call_tokens
 from bex.ensemble import infer_ensemble
 
+IMPORT_PATTERNS = [
+    re.compile(r"^\s*import\s+"),
+    re.compile(r"^\s*from\s+"),
+    re.compile(r"^\s*require\s+"),
+    re.compile(r"^\s*require_relative\s+"),
+    re.compile(r"^\s*#\s*include\s+"),
+    re.compile(r"^\s*use\s+"),
+    re.compile(r"^\s*include\s+"),
+]
+
 SUPPORTED_EXTENSIONS = {
     ".py", ".js", ".ts", ".kt", ".rb", ".go", ".rs", ".java", ".c", ".cpp",
 }
@@ -46,6 +58,30 @@ def _match_glob(filepath, pattern):
     return spec.match_file(filepath)
 
 
+def _extract_imports(file_paths):
+    """Extract unique import lines from source files.
+
+    Scans top 200 lines of each file for common import patterns
+    across all 10 supported languages. Deduplicates across files.
+    """
+    seen = set()
+    result = []
+    for fp in sorted(file_paths):
+        try:
+            with open(fp) as f:
+                for i, line in enumerate(f):
+                    if i >= 200:
+                        break
+                    stripped = line.strip()
+                    if any(p.match(stripped) for p in IMPORT_PATTERNS):
+                        if stripped not in seen:
+                            seen.add(stripped)
+                            result.append(stripped)
+        except OSError:
+            continue
+    return result
+
+
 def scan_directory(dir_path, gitignore_spec=None):
     """Walk dir_path, return dict mapping extension → [file paths].
 
@@ -165,15 +201,18 @@ def analyze_clusters(file_paths, extension, min_coverage=0.2, prefer=None, kmax=
     """Run full pipeline with clustering: preprocess → cluster → per-cluster infer.
 
     Returns:
-        list of (label, ensemble_result_dict, method_count) tuples.
+        list of (label, ensemble_result_dict, method_count, meta) tuples.
+        meta = {"files": set(paths), "imports": [sorted_import_lines]}.
     """
     sequences = []
+    seq_files = []
     for fp in file_paths:
         with open(fp) as f:
             code = f.read()
         for method_seq in preprocess_by_method(fp, code):
             if method_seq:
                 sequences.append(method_seq)
+                seq_files.append(fp)
 
     if not sequences:
         return []
@@ -183,9 +222,15 @@ def analyze_clusters(file_paths, extension, min_coverage=0.2, prefer=None, kmax=
 
     results = []
     for label, cluster_seqs in clusters:
+        cluster_fps = set()
+        for seq in cluster_seqs:
+            idx = next(i for i, s in enumerate(sequences) if s is seq)
+            cluster_fps.add(seq_files[idx])
+        imports = _extract_imports(cluster_fps)
         symbol_seqs = [[text for _, text, _ in seq] for seq in cluster_seqs]
         result = infer_ensemble(symbol_seqs, kmax=kmax, N=N, prefer=prefer)
-        results.append((label, result, len(cluster_seqs)))
+        meta = {"files": cluster_fps, "imports": imports}
+        results.append((label, result, len(cluster_seqs), meta))
 
     return results
 
@@ -258,7 +303,7 @@ def analyze_directory(
             )
         else:
             r = infer(files, ext, min_coverage=min_coverage, prefer=prefer, kmax=kmax)
-            results[ext] = [("(all methods)", r, 0)]
+            results[ext] = [("(all methods)", r, 0, {"files": set(files), "imports": _extract_imports(files)})]
     return results
 
 
@@ -292,9 +337,40 @@ def _parse_args(argv=None):
         "--ngram-size", type=int, default=3,
         help="N-gram length for clustering (default: 3)",
     )
+    parser.add_argument(
+        "--format", choices=["text", "json"], default="text",
+        help="Output format (default: text)",
+    )
+    parser.add_argument(
+        "--json", action="store_true", dest="json_flag",
+        help="Shortcut for --format json",
+    )
     return parser.parse_args(argv)
 
 
+def _build_json_output(results):
+    """Convert results dict to a compact JSON structure for prompt injection."""
+    output = []
+    for ext, clusters in results.items():
+        lang = {"language": ext, "conventions": []}
+        total_methods = 0
+        for label, result, count, meta in clusters:
+            total_methods += count
+            entry = {
+                "label": label,
+                "method_count": count,
+            }
+            if result and result.get("best"):
+                entry["algorithm"] = result["best"]["algorithm"]
+                entry["grammar"] = result["best"]["grammar"]
+                entry["mdl_score"] = round(result["best"]["mdl_score"], 1)
+            entry["imports"] = meta.get("imports", [])
+            lang["conventions"].append(entry)
+        lang["total_methods"] = total_methods
+        output.append(lang)
+    return json.dumps(output, indent=2)
+
+
 def main():
     args = _parse_args()
     results = analyze_directory(
@@ -304,9 +380,14 @@ def main():
         kmax=args.kmax,
         include=args.include,
     )
+
+    if args.json_flag or args.format == "json":
+        print(_build_json_output(results))
+        return
+
     for ext, clusters in results.items():
         print(f"\n{ext}:")
-        for label, result, count in clusters:
+        for label, result, count, meta in clusters:
             if result and result.get("best"):
                 best = result["best"]
                 print(f"  ╰─ {label} ({count} methods)")
@@ -315,6 +396,12 @@ def main():
                 print(f"     MDL:       {best['mdl_score']}")
             else:
                 print(f"  ╰─ {label} ({count} methods) — no grammar")
+            imps = meta.get("imports", [])
+            if imps:
+                joined = "  |  ".join(imps[:6])
+                print(f"     Imports:  {joined}")
+                if len(imps) > 6:
+                    print(f"               ... and {len(imps) - 6} more")
 
 
 if __name__ == "__main__":

tests/test_analyze.py

@@ -138,7 +138,7 @@ def test_analyze_directory_include_glob():
         results = analyze_directory(td, include="**/src/main/**")
         assert ".py" in results
         assert len(results[".py"]) >= 1
-        for label, r, count in results[".py"]:
+        for label, r, count, meta in results[".py"]:
             if r and r.get("best"):
                 assert r["best"]["grammar"] is not None
     print("  PASS test_analyze_directory_include_glob")