resolver.py

#!/usr/bin/env python3
"""
resolver.py — resolves Python imports to either:
  - a stub file in stubs/<name>.py  (for C extensions / stdlib)
  - the module's actual .py source  (auto-transpiled)

C extensions with no stub raise TranspileError with a clear message.
Relative imports (from . import X) are skipped by the resolver — they reference
files already loaded as part of the same package.
"""
import ast
import importlib.util
from pathlib import Path
from typing import List, Tuple, Set

from core import TranspileError

STUBS_DIR = Path(__file__).parent / "stubs"


class ImportResolver:
    """
    Recursively resolves all imports reachable from a set of modules.
    Returns only modules that were not already known.

    For packages (directories with __init__.py) every .py file in the
    package root is loaded so that relative intra-package imports resolve.
    """

    def __init__(self, known_names: Set[str]):
        self._known: Set[str] = set(known_names)

    def collect(self, modules: List[ast.Module]) -> List[Tuple[ast.Module, str]]:
        """
        Scan imports in `modules`, resolve unknown ones, recurse into them.
        Returns list of (ast.Module, module_name) pairs to add to transpilation.
        """
        result: List[Tuple[ast.Module, str]] = []
        pending = list(modules)
        while pending:
            mod = pending.pop()
            for new_tree, new_name in self._resolve_mod_imports(mod):
                result.append((new_tree, new_name))
                pending.append(new_tree)
        return result

    # ------------------------------------------------------------------

    def _resolve_mod_imports(self, mod: ast.Module) -> List[Tuple[ast.Module, str]]:
        seen_here: Set[str] = set()
        result: List[Tuple[ast.Module, str]] = []
        for node in mod.body:
            names: List[str] = []
            if isinstance(node, ast.Import):
                names = [alias.name.split(".")[0] for alias in node.names]
            elif isinstance(node, ast.ImportFrom) and node.module and node.level == 0:
                # Only absolute imports; relative imports (level > 0) are
                # intra-package and handled by loading the whole package.
                names = [node.module.split(".")[0]]
            for name in names:
                if name not in self._known and name not in seen_here:
                    seen_here.add(name)
                    for r_tree, r_name in self._resolve_one(name):
                        self._known.add(r_name)
                        result.append((r_tree, r_name))
        return result

    def _resolve_one(self, name: str) -> List[Tuple[ast.Module, str]]:
        # 1. Stub takes priority (C extension or hand-written mapping)
        stub = STUBS_DIR / f"{name}.py"
        if stub.exists():
            src = stub.read_text(encoding="utf-8")
            tree = ast.parse(src, filename=str(stub))
            return [(tree, name)]

        # 2. Try to locate Python source via importlib
        try:
            spec = importlib.util.find_spec(name)
        except (ModuleNotFoundError, ValueError):
            spec = None

        if spec is None:
            raise TranspileError(f"Module '{name}' not found")

        # Plain .py file
        if spec.origin and spec.origin.endswith(".py"):
            src = Path(spec.origin).read_text(encoding="utf-8")
            tree = ast.parse(src, filename=spec.origin)
            return [(tree, name)]

        # Package: load __init__.py + every submodule .py in the package root.
        # This ensures relative intra-package imports (from .models import ...)
        # are satisfied because those files are already in the module list.
        if spec.submodule_search_locations:
            results: List[Tuple[ast.Module, str]] = []
            for loc in spec.submodule_search_locations:
                loc_path = Path(loc)
                for py_file in sorted(loc_path.glob("*.py")):
                    try:
                        src = py_file.read_text(encoding="utf-8")
                        tree = ast.parse(src, filename=str(py_file))
                    except (OSError, SyntaxError):
                        continue
                    mod_name = name if py_file.stem == "__init__" else f"{name}.{py_file.stem}"
                    results.append((tree, mod_name))
                break  # only first search location
            if results:
                return results

        # C extension with no stub
        raise TranspileError(
            f"Module '{name}' is a C extension (no .py source found) — "
            f"create stubs/{name}.py to support it"
        )