testsAndMisc/python_pkg/code_tutor/_analyzer.py

"""Language-agnostic code item extraction from source files.

Supports Python via the ``ast`` module (exact line ranges) and other languages
via a regex heuristic.  Binary files and vendored directories are skipped.
"""

from __future__ import annotations

import ast
from dataclasses import dataclass, field
import hashlib
from pathlib import Path
import re

_SKIP_DIRS: frozenset[str] = frozenset(
    {
        "third_party",
        ".venv",
        "node_modules",
        "__pycache__",
        ".git",
        "dist",
        "build",
    }
)
_SKIP_SUFFIXES: frozenset[str] = frozenset(
    {
        ".geojson",
        ".png",
        ".jpg",
        ".jpeg",
        ".gif",
        ".ico",
        ".svg",
        ".webp",
        ".pdf",
        ".zip",
        ".tar",
        ".gz",
        ".bin",
    }
)
_OTHER_LANGS: frozenset[str] = frozenset(
    {".js", ".ts", ".go", ".rs", ".c", ".cpp", ".dart"}
)

# Matches leading keyword(s) followed by an identifier and opening paren.
_FUNCTION_RE: re.Pattern[str] = re.compile(
    r"^[ \t]*(?:(?:pub(?:lic)?|priv(?:ate)?|prot(?:ected)?|static|async)\s+)*"
    r"(?:def|function|fn|func|void|int|float|bool|string)\s+(\w+)\s*\(",
    re.MULTILINE,
)


@dataclass
class CodeItem:
    """A single extractable code unit (function or method).

    Attributes:
        id: Dotted identifier, e.g. ``module.submodule.function_name``.
        file: Relative path from the codebase root.
        type: ``"function"`` or ``"async_function"``.
        name: Bare function or method name.
        start_line: 1-based first line of the definition.
        end_line: 1-based last line of the definition.
        class_name: Enclosing class name, or ``""`` for module-level functions.
        depends_on: IDs of items that must be understood first.
    """

    id: str
    file: str
    type: str
    name: str
    start_line: int
    end_line: int
    class_name: str = ""
    depends_on: list[str] = field(default_factory=list)


def _is_binary(path: Path) -> bool:
    """Return True if *path* appears to be a binary file (null-byte heuristic).

    Uses the same heuristic as ``git`` and GNU ``grep``: a file is binary if
    its first 512 bytes contain a null byte.  Raises ``OSError`` on read
    failure so callers can handle unreadable files separately.

    Args:
        path: File to inspect.

    Returns:
        True when the first 512 bytes contain a null byte; False otherwise.

    Raises:
        OSError: When the file cannot be read.
    """
    chunk = path.read_bytes()[:512]
    return b"\x00" in chunk


def _should_skip(rel: str) -> bool:
    """Return True if this relative path should be excluded from analysis.

    Args:
        rel: Path relative to the codebase root.

    Returns:
        True when any path component is in ``_SKIP_DIRS`` or the file
        extension is in ``_SKIP_SUFFIXES``.
    """
    parts = Path(rel).parts
    if any(p in _SKIP_DIRS for p in parts):
        return True
    return Path(rel).suffix.lower() in _SKIP_SUFFIXES


def _make_id(rel_path: str, name: str) -> str:
    """Build a dotted item identifier from a relative path and function name.

    Args:
        rel_path: Path relative to the codebase root.
        name: Function or method name.

    Returns:
        Dotted identifier like ``module.submodule.function_name``.
    """
    without_ext = Path(rel_path).with_suffix("")
    return ".".join(without_ext.parts) + "." + name


class _FunctionVisitor(ast.NodeVisitor):
    """AST visitor that extracts function/method items with class context.

    Args:
        rel_path: Relative path of the file being visited.
    """

    def __init__(self, rel_path: str) -> None:
        """Initialise with file path and empty state."""
        self._rel_path = rel_path
        self._class_stack: list[str] = []
        self.items: list[CodeItem] = []

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        """Push class name onto stack, recurse, then pop."""
        self._class_stack.append(node.name)
        self.generic_visit(node)
        self._class_stack.pop()

    def _add(self, node: ast.FunctionDef | ast.AsyncFunctionDef, kind: str) -> None:
        end = getattr(node, "end_lineno", node.lineno)
        self.items.append(
            CodeItem(
                id=_make_id(self._rel_path, node.name),
                file=self._rel_path,
                type=kind,
                name=node.name,
                start_line=node.lineno,
                end_line=end,
                class_name=self._class_stack[-1] if self._class_stack else "",
            )
        )

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        """Record function and recurse into body."""
        self._add(node, "function")
        self.generic_visit(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
        """Record async function and recurse into body."""
        self._add(node, "async_function")
        self.generic_visit(node)


def _extract_python(path: Path, rel_path: str) -> list[CodeItem]:
    """Extract function/method items from a Python file using the AST.

    Args:
        path: Absolute path to the Python file.
        rel_path: Relative path used to build item IDs.

    Returns:
        List of ``CodeItem`` instances, one per function definition found.
        Returns an empty list on ``SyntaxError``.
    """
    try:
        source = path.read_text(encoding="utf-8", errors="replace")
        tree = ast.parse(source)
    except SyntaxError:
        return []

    visitor = _FunctionVisitor(rel_path)
    visitor.visit(tree)
    return visitor.items


def _extract_other(path: Path, rel_path: str) -> list[CodeItem]:
    """Extract function-like items from non-Python files via regex.

    Scans each line for a function-definition keyword followed by an
    identifier and opening parenthesis.  End line is estimated at
    ``start + 40`` (clamped to file length).

    Args:
        path: Absolute path to the source file.
        rel_path: Relative path used to build item IDs.

    Returns:
        List of ``CodeItem`` instances.  Returns an empty list on ``OSError``.
    """
    try:
        lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
    except OSError:
        return []

    items: list[CodeItem] = []
    for lineno, line in enumerate(lines, start=1):
        match = _FUNCTION_RE.match(line)
        if match:
            name = match.group(1)
            end = min(lineno + 40, len(lines))
            items.append(
                CodeItem(
                    id=_make_id(rel_path, name),
                    file=rel_path,
                    type="function",
                    name=name,
                    start_line=lineno,
                    end_line=end,
                )
            )
    return items


def extract_items(codebase: Path) -> list[CodeItem]:
    """Walk *codebase* and return all extractable code items.

    Args:
        codebase: Root directory to analyse.

    Returns:
        Flat list of ``CodeItem`` instances, unsorted.
    """
    items: list[CodeItem] = []
    for path in sorted(codebase.rglob("*")):
        if not path.is_file():
            continue
        rel = str(path.relative_to(codebase))
        if _should_skip(rel):
            continue
        try:
            binary = _is_binary(path)
        except OSError:
            continue
        if binary:
            continue
        if path.suffix == ".py":
            items.extend(_extract_python(path, rel))
        elif path.suffix in _OTHER_LANGS:
            items.extend(_extract_other(path, rel))
    return items


def get_python_files(codebase: Path) -> list[Path]:
    """Return all non-skipped Python files under *codebase*.

    Args:
        codebase: Root directory to search.

    Returns:
        Sorted list of absolute Python file paths.
    """
    return [
        path
        for path in sorted(codebase.rglob("*.py"))
        if path.is_file() and not _should_skip(str(path.relative_to(codebase)))
    ]


def _build_dotted_map(codebase: Path, all_files: list[Path]) -> dict[Path, str]:
    """Build a mapping from file to its dotted module path relative to *codebase*.

    Args:
        codebase: Root directory used for relative path computation.
        all_files: All Python files within the codebase.

    Returns:
        Dict mapping each file in *all_files* to its dotted module string.
        Files that cannot be made relative to *codebase* are excluded.
    """
    result: dict[Path, str] = {}
    for f in all_files:
        try:
            rel = f.relative_to(codebase)
        except ValueError:
            continue
        result[f] = ".".join(rel.with_suffix("").parts)
    return result


def _match_deps(
    imported: set[str],
    dotted_map: dict[Path, str],
    path: Path,
) -> set[Path]:
    """Return the subset of *dotted_map* files that *path* appears to import.

    A file ``f`` is included when any imported module name equals, contains,
    or is contained by ``f``'s dotted path.

    Args:
        imported: Set of module name strings collected from import statements.
        dotted_map: Mapping of file to dotted module path.
        path: The file being analysed (excluded from its own dependency set).

    Returns:
        Set of dependency files.
    """
    deps: set[Path] = set()
    for imp in imported:
        for f, dotted in dotted_map.items():
            if f != path and (
                imp == dotted
                or imp.endswith("." + dotted)
                or dotted.endswith("." + imp)
            ):
                deps.add(f)
    return deps


def get_file_imports(path: Path, codebase: Path, all_files: list[Path]) -> set[Path]:
    """Parse Python imports and return the codebase files this file depends on.

    Uses a suffix-match heuristic: if an imported module name ends with (or
    matches) the dotted path of a file in *all_files*, that file is treated as
    a dependency.

    Args:
        path: The Python file to analyse.
        codebase: Root directory used for relative path computation.
        all_files: All Python files within the codebase.

    Returns:
        Subset of *all_files* that *path* imports from.
    """
    try:
        source = path.read_text(encoding="utf-8", errors="replace")
        tree = ast.parse(source)
    except SyntaxError:
        return set()

    imported: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                imported.add(alias.name)
        elif isinstance(node, ast.ImportFrom) and node.module:
            imported.add(node.module)

    dotted_map = _build_dotted_map(codebase, all_files)
    return _match_deps(imported, dotted_map, path)


def codebase_fingerprint(codebase: Path) -> str:
    """Compute a short fingerprint of all source files under *codebase*.

    Uses each file's mtime and size rather than content, so it is fast
    even on large repositories.  Returns a 16-character hex string that
    changes whenever any tracked file is added, removed, or modified.

    Args:
        codebase: Root directory of the codebase.

    Returns:
        16-character hex fingerprint string.
    """
    parts: list[str] = []
    for path in sorted(codebase.rglob("*")):
        if not path.is_file():
            continue
        rel = str(path.relative_to(codebase))
        if _should_skip(rel):
            continue
        try:
            s = path.stat()
            parts.append(f"{rel}:{s.st_mtime_ns}:{s.st_size}")
        except OSError:
            pass
    return hashlib.sha256("\n".join(parts).encode()).hexdigest()[:16]