#!/usr/bin/env python3 """ fix_imports.py — Make intra-package imports in a src/ layout relative. Usage: python tools/fix_imports.py \ --root packages/vinefeeder/src \ --pkg-name vinefeeder \ [--apply] By default it prints a plan (dry-run). Add --apply to rewrite files. Backups (*.bak) are created when applying changes. """ from __future__ import annotations import argparse import ast import difflib import shutil from pathlib import Path from typing import Iterable, List, Tuple def find_package_modules(pkg_dir: Path) -> set[str]: """ Return a set of top-level module/package names inside pkg_dir. Example: { 'pretty', 'parsing_utils', 'batchloader', 'services', ... } """ names: set[str] = set() for p in pkg_dir.iterdir(): if p.name.startswith("_"): continue if p.is_dir() and (p / "__init__.py").exists(): names.add(p.name) elif p.suffix == ".py": names.add(p.stem) return names def rewrite_imports( src: str, file_path: Path, pkg_name: str, local_modules: set[str], ) -> Tuple[str, List[str]]: """ Return (new_source, notes). Keeps formatting line-by-line as much as possible. We only rewrite simple import statements using AST nodes to decide the intent, then patch lines textually. """ notes: List[str] = [] try: tree = ast.parse(src) except SyntaxError as e: notes.append(f"SKIP (syntax error): {e}") return src, notes # Collect edits as (lineno-1, new_line) for lines we fully replace edits: dict[int, str] = {} # We’ll need original lines to reconstruct edits lines = src.splitlines(keepends=False) for node in tree.body: # Handle: import X [, Y] [as Z] if isinstance(node, ast.Import): # If *all* top-level names refer to local modules, we can rewrite: # import pretty -> from . import pretty # import pretty as p -> from . import pretty as p # import pretty, parsing_utils -> from . import pretty, parsing_utils local_aliases = [] non_local = False for alias in node.names: top = alias.name.split(".")[0] if top in local_modules or top == pkg_name: local_aliases.append(alias) else: non_local = True break if non_local or not local_aliases: continue # leave it # Build a single "from . import ..." line preserving aliases parts = [] for alias in node.names: top = alias.name.split(".")[0] # Only rewrite locals, leave non-locals untouched (rare mixed case) if top in local_modules or top == pkg_name: if alias.asname: parts.append(f"{alias.name} as {alias.asname}") else: parts.append(alias.name) else: # If we ever hit mixed, bail out and don't rewrite parts = [] break if not parts: continue new_line = f"from . import {', '.join(parts)}" # Replace the entire original line (best-effort: single-line import) lineno = node.lineno - 1 old = lines[lineno].strip() edits[lineno] = new_line notes.append(f"import→relative: `{old}` -> `{new_line}`") # Handle: from X import Y elif isinstance(node, ast.ImportFrom): # Already relative? if node.level and node.level > 0: continue if node.module is None: continue top = node.module.split(".")[0] # Case A: from vinefeeder.something import thing if top == pkg_name: # Convert to from .something import thing (preserve the tail) tail = node.module.split(".", 1)[1] if "." in node.module else "" dot_path = f".{tail}" if tail else "." names = [] for alias in node.names: if alias.asname: names.append(f"{alias.name} as {alias.asname}") else: names.append(alias.name) new_line = f"from {dot_path} import {', '.join(names)}" lineno = node.lineno - 1 old = lines[lineno].strip() edits[lineno] = new_line notes.append(f"abs→relative: `{old}` -> `{new_line}`") continue # Case B: from sibling import thing (e.g., from pretty import foo) if top in local_modules: # Keep subpath if present (e.g., services.util) tail = node.module new_line = f"from .{tail} import " + ", ".join( f"{n.name} as {n.asname}" if n.asname else n.name for n in node.names ) lineno = node.lineno - 1 old = lines[lineno].strip() edits[lineno] = new_line notes.append(f"sibling→relative: `{old}` -> `{new_line}`") if not edits: return src, notes # Apply edits (line-level replacement) new_lines = [] for i, line in enumerate(lines): new_lines.append(edits.get(i, line)) new_src = "\n".join(new_lines) + ("\n" if src.endswith("\n") else "") return new_src, notes def iter_python_files(pkg_dir: Path) -> Iterable[Path]: for p in pkg_dir.rglob("*.py"): yield p def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--root", required=True, help="Path to the src/ root (e.g., packages/vinefeeder/src)") ap.add_argument("--pkg-name", required=True, help="Top-level package name (e.g., vinefeeder)") ap.add_argument("--apply", action="store_true", help="Rewrite files in place (creates .bak backups)") args = ap.parse_args() src_root = Path(args.root).resolve() pkg_dir = (src_root / args.pkg_name).resolve() if not pkg_dir.exists(): print(f"[error] Package dir not found: {pkg_dir}") return 2 local_modules = find_package_modules(pkg_dir) print(f"[info] Package: {args.pkg_name}") print(f"[info] Package dir: {pkg_dir}") print(f"[info] Local modules found: {sorted(local_modules)}") print() total_changes = 0 for py in iter_python_files(pkg_dir): old = py.read_text(encoding="utf-8") new, notes = rewrite_imports(old, py, args.pkg_name, local_modules) if not notes: continue total_changes += 1 print(f"--- {py}") for n in notes: print(" -", n) if args.apply and new != old: bak = py.with_suffix(py.suffix + ".bak") shutil.copy2(py, bak) py.write_text(new, encoding="utf-8") # Show a small unified diff for context diff = difflib.unified_diff( old.splitlines(), new.splitlines(), fromfile=str(py), tofile=str(py), lineterm="" ) print("\n".join(diff)) print() if total_changes == 0: print("[info] No imports to rewrite.") else: print(f"[done] Files with changes: {total_changes} {'(dry-run)' if not args.apply else ''}") return 0 if __name__ == "__main__": raise SystemExit(main())