scripts/p3i_hoist_lazy_imports.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""P3-I: hoist function-scoped ``from butler.*`` to module top where safe."""
from __future__ import annotations
import argparse
import ast
import pathlib
import subprocess
import sys
import textwrap
# Proven import cycles — keep lazy in these (file_posix, module).
CYCLE_KEEP: set[tuple[str, str]] = set()
MIN_COUNT = 8
def _import_key(node: ast.ImportFrom) -> tuple:
names = tuple(sorted((a.name, a.asname) for a in (node.names or [])))
return (node.level or 0, node.module or "", names)
def _import_stmt(node: ast.ImportFrom) -> str:
return ast.unparse(node).strip()
def _collect_imports(tree: ast.AST) -> tuple[set[tuple], list[tuple[ast.ImportFrom, int]]]:
"""Return (module_level_keys, lazy_nodes_with_depth)."""
module_keys: set[tuple] = set()
lazy: list[tuple[ast.ImportFrom, int]] = []
class V(ast.NodeVisitor):
def __init__(self) -> None:
self.depth = 0
self.in_type_checking = False
def visit_If(self, node: ast.If) -> None:
guard = False
if isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING":
guard = True
elif (
isinstance(node.test, ast.Attribute)
and isinstance(node.test.value, ast.Name)
and node.test.value.id == "typing"
and node.test.attr == "TYPE_CHECKING"
):
guard = True
prev = self.in_type_checking
if guard:
self.in_type_checking = True
self.generic_visit(node)
self.in_type_checking = prev
def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
self.depth += 1
self.generic_visit(node)
self.depth -= 1
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
self.depth += 1
self.generic_visit(node)
self.depth -= 1
def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
mod = node.module or ""
if not mod.startswith("butler"):
return
key = _import_key(node)
if self.depth == 0 and not self.in_type_checking:
module_keys.add(key)
elif self.depth > 0:
lazy.append((node, self.depth))
V().visit(tree)
return module_keys, lazy
def _insert_after_imports(lines: list[str], new_stmts: list[str]) -> list[str]:
"""Insert new import lines after the last contiguous top-level import block."""
insert_at = 0
i = 0
n = len(lines)
# Skip optional module docstring (triple-quoted string at file start).
if i < n:
stripped0 = lines[i].strip()
if stripped0.startswith('"""') or stripped0.startswith("'''"):
quote = stripped0[:3]
if stripped0.count(quote) >= 2 and len(stripped0) > 6:
insert_at = i + 1
i += 1
else:
i += 1
while i < n and quote not in lines[i]:
i += 1
if i < n:
insert_at = i + 1
i += 1
while i < n:
line = lines[i]
stripped = line.strip()
if stripped.startswith("from __future__"):
insert_at = i + 1
i += 1
continue
if stripped.startswith("import ") or stripped.startswith("from "):
insert_at = i + 1
# skip multiline import paren blocks
if "(" in stripped and ")" not in stripped:
i += 1
while i < n and ")" not in lines[i]:
insert_at = i + 1
i += 1
if i < n:
insert_at = i + 1
i += 1
continue
if stripped == "" or stripped.startswith("#"):
i += 1
continue
break
out = lines[:insert_at]
if out and out[-1].strip():
out.append("")
for stmt in new_stmts:
out.append(stmt if stmt.endswith("\n") else stmt + "\n")
if insert_at < n and (not out or out[-1].strip()):
pass
elif insert_at < n:
out.append("")
out.extend(lines[insert_at:])
return out
def _remove_node_lines(source: str, node: ast.AST) -> str:
lines = source.splitlines(keepends=True)
start = node.lineno - 1
end = getattr(node, "end_lineno", node.lineno) - 1
# drop trailing blank line if import was alone in block
new_lines = lines[:start] + lines[end + 1 :]
return "".join(new_lines)
def process_file(path: pathlib.Path, *, dry_run: bool = False) -> dict:
rel = path.as_posix()
source = path.read_text(encoding="utf-8")
try:
tree = ast.parse(source)
except SyntaxError as e:
return {"path": rel, "error": str(e), "removed": 0, "hoisted": 0}
module_keys, lazy_nodes = _collect_imports(tree)
if not lazy_nodes:
return {"path": rel, "removed": 0, "hoisted": 0, "kept_cycle": 0}
to_remove: list[ast.ImportFrom] = []
to_hoist: list[ast.ImportFrom] = []
kept_cycle = 0
for node, _depth in lazy_nodes:
mod = node.module or ""
key = _import_key(node)
if (rel, mod) in CYCLE_KEEP:
kept_cycle += 1
continue
if key in module_keys:
to_remove.append(node)
else:
to_hoist.append(node)
module_keys.add(key)
if not to_remove and not to_hoist:
return {"path": rel, "removed": 0, "hoisted": 0, "kept_cycle": kept_cycle}
# Remove bottom-up to preserve line numbers
all_remove = sorted(to_remove + to_hoist, key=lambda n: n.lineno, reverse=True)
new_source = source
for node in all_remove:
new_source = _remove_node_lines(new_source, node)
hoist_stmts = []
seen_hoist: set[tuple] = set()
for node in to_hoist:
key = _import_key(node)
if key in seen_hoist:
continue
seen_hoist.add(key)
hoist_stmts.append(_import_stmt(node))
if hoist_stmts:
lines = new_source.splitlines(keepends=True)
if lines and not lines[-1].endswith("\n"):
lines[-1] += "\n"
new_source = "".join(_insert_after_imports(lines, hoist_stmts))
if dry_run:
return {
"path": rel,
"removed": len(to_remove),
"hoisted": len(seen_hoist),
"kept_cycle": kept_cycle,
}
path.write_text(new_source, encoding="utf-8")
return {
"path": rel,
"removed": len(to_remove),
"hoisted": len(seen_hoist),
"kept_cycle": kept_cycle,
}
def verify_import() -> bool:
r = subprocess.run(
[sys.executable, "-c", "import butler; import butler.main"],
cwd=pathlib.Path("."),
env={**dict(**__import__("os").environ), "PYTHONPATH": "."},
capture_output=True,
text=True,
)
return r.returncode == 0
def count_lazy() -> int:
sys.path.insert(0, ".")
from butler.ops.lazy_import_budget import count_lazy_butler_imports
return count_lazy_butler_imports()
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--min", type=int, default=MIN_COUNT)
parser.add_argument("--file", action="append", dest="files")
args = parser.parse_args()
sys.path.insert(0, ".")
from butler.ops.lazy_import_budget import lazy_import_counts_by_file
if args.files:
targets = [pathlib.Path(f) for f in args.files]
else:
by_file = lazy_import_counts_by_file()
targets = [
pathlib.Path(p) for p, c in sorted(by_file.items(), key=lambda x: -x[1]) if c >= args.min
]
start = count_lazy()
print(f"START: {start}")
changed: list[dict] = []
batch_size = 10
for i, path in enumerate(targets):
result = process_file(path, dry_run=args.dry_run)
if result.get("removed") or result.get("hoisted"):
changed.append(result)
print(
f" {result['path']}: removed_dup={result.get('removed',0)} "
f"hoisted={result.get('hoisted',0)} kept_cycle={result.get('kept_cycle',0)}"
)
if not args.dry_run and (i + 1) % batch_size == 0:
if not verify_import():
print("IMPORT FAILED after batch — stop", file=sys.stderr)
print(f"COUNT: {count_lazy()}")
return 1
print(f" batch {(i+1)//batch_size} ok, count={count_lazy()}")
if not args.dry_run:
if not verify_import():
print("IMPORT FAILED at end", file=sys.stderr)
return 1
end = count_lazy()
print(f"END: {end} (delta {start - end})")
print(f"FILES CHANGED: {len(changed)}")
return 0
if __name__ == "__main__":
raise SystemExit(main())