`ast.NodeVisitor` for import tracking by ShaharNaveh · Pull Request #7229 · RustPython/RustPython
def _extract_top_level_code(content: str) -> str: """Extract only top-level code from Python content for faster parsing.""" def_idx = content.find("\ndef ") class_idx = content.find("\nclass ") class ImportVisitor(ast.NodeVisitor): def __init__(self) -> None: self.test_imports = set() self.lib_imports = set()
indices = [i for i in (def_idx, class_idx) if i != -1] if indices: content = content[: min(indices)] return content.rstrip("\n") def add_import(self, name: str) -> None: """ Add an `import` to its correct slot (`test_imports` or `lib_imports`).
Parameters ---------- name : str Module name. """ if name.startswith("test.support"): return
_FROM_TEST_IMPORT_RE = re.compile(r"^from test import (.+)", re.MULTILINE) _FROM_TEST_DOT_RE = re.compile(r"^from test\.(\w+)", re.MULTILINE) _IMPORT_TEST_DOT_RE = re.compile(r"^import test\.(\w+)", re.MULTILINE) real_name = name.split(".", 1)[-1] if name.startswith("test."): self.test_imports.add(real_name) else: self.lib_imports.add(real_name)
def visit_Import(self, node): for alias in node.names: self.add_import(alias.name)
def parse_test_imports(content: str) -> set[str]: """Parse test file content and extract test package dependencies.""" content = _extract_top_level_code(content) imports = set() def visit_ImportFrom(self, node): try: module = node.module except AttributeError: # Ignore `from . import my_internal_module` return
for match in _FROM_TEST_IMPORT_RE.finditer(content): import_list = match.group(1) for part in import_list.split(","): name = part.split()[0].strip() if name and name not in ("support", "__init__"): imports.add(name) for name in node.names: self.add_import(f"{module}.{name}")
for match in _FROM_TEST_DOT_RE.finditer(content): dep = match.group(1) if dep not in ("support", "__init__"): imports.add(dep) def visit_Call(self, node) -> None: """ In test files, there's sometimes use of:
for match in _IMPORT_TEST_DOT_RE.finditer(content): dep = match.group(1) if dep not in ("support", "__init__"): imports.add(dep) ```python import test.support from test.support import script_helper
return imports script = support.findfile("_test_atexit.py") script_helper.run_test_script(script) ```
This imports "_test_atexit.py" but does not show as an import node. """ func = node.func if not isinstance(func, ast.Attribute): return
_IMPORT_RE = re.compile(r"^import\s+(\w[\w.]*)", re.MULTILINE) _FROM_IMPORT_RE = re.compile(r"^from\s+(\w[\w.]*)\s+import", re.MULTILINE) value = func.value if not isinstance(value, ast.Name): return
if (value.id != "support") or (func.attr != "findfile"): return
def parse_lib_imports(content: str) -> set[str]: """Parse library file and extract all imported module names.""" imports = set() arg = node.args[0] if not isinstance(arg, ast.Constant): return
target = arg.value if not target.endswith(".py"): return
for match in _IMPORT_RE.finditer(content): imports.add(match.group(1)) target = target.removesuffix(".py") self.add_import(f"test.{target}")
for match in _FROM_IMPORT_RE.finditer(content): imports.add(match.group(1))
return imports def parse_test_imports(content: str) -> set[str]: """Parse test file content and extract test package dependencies.""" if not (tree := safe_parse_ast(content)): return set()
visitor = ImportVisitor() visitor.visit(tree) return visitor.test_imports
def parse_lib_imports(content: str) -> set[str]: """Parse library file and extract all imported module names.""" if not (tree := safe_parse_ast(content)): return set()
visitor = ImportVisitor() visitor.visit(tree) return visitor.lib_imports
# === TODO marker utilities ===
def count_rustpython_todo(content: str) -> int: """Count lines containing RustPython TODO markers.""" return sum(1 for line in content.splitlines() if TODO_MARKER in line) return content.count(TODO_MARKER)
def count_todo_in_path(path: pathlib.Path) -> int:
total = 0 for _, content in read_python_files(path): total += count_rustpython_todo(content) return total return sum(count_rustpython_todo(content) for _, content in read_python_files(path))
# === Test utilities ===