orchestrata_tools.parser.source_parser

Reference › orchestrata_tools › parser › source_parser

SourceASTParser parses Python source files directly into ASTCodebase models using only the standard ast module. It lets orchestrata_tools analyse any Python project without first running the serializer. It is used by the Jinja and stub test generators' tests (tests/test_orchestrata_tools_jinja_gen.py, tests/test_orchestrata_tools_stub_generator.py) and covered by tests/test_orchestrata_tools_source_parser_coverage.py and tests/test_source_parser_attrs.py. The loader for existing snapshots is json_loader; the models are in ast_models.

parse_directory walks a directory, optionally skipping excluded names, and returns an ASTCodebase. parse_file parses one file into a ModuleInfo, with the module path computed relative to an optional root. Per-node helpers build classes, methods, functions, parameters, imports and module variables, and _extract_calls collects the call names inside a body. Two constant sets decide flags from base-class names: the final dotted segment of a base marks an Enum (Enum, IntEnum, StrEnum, IntFlag, Flag, ReprEnum) or an abstract class (ABC, ABCMeta), in addition to an abstract method.

At a glance

Class diagram

Dependency diagram

API

Classes

SourceASTParser

Line 36

Parse Python source files into ASTCodebase without external dependencies.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_tools/parser/source_parser.py

"""Parse Python source files directly into ASTCodebase models.

Uses only Python's built-in `ast` module. No third-party dependencies.
This allows orchestrata_tools to analyze any Python project without
requiring the serialize_codebase_ast.py script to be run first.
"""

from __future__ import annotations

import ast
import sys
from pathlib import Path
from typing import Any

from orchestrata_tools.models.ast_models import (
    ASTCodebase,
    ClassInfo,
    ClassVarInfo,
    FunctionInfo,
    ImportInfo,
    MethodInfo,
    ModuleInfo,
    ParamInfo,
)

# Final segment (after the last ".") of base-class expressions that mark a
# class as an Enum, e.g. "Enum", "IntEnum", "enum.Enum", "MyProject.StrEnum".
_ENUM_BASE_FINAL_NAMES = frozenset(
    {"Enum", "IntEnum", "StrEnum", "IntFlag", "Flag", "ReprEnum"}
)
# Final segment of base-class expressions that mark a class as abstract via
# inheritance (in addition to containing an @abstractmethod).
_ABSTRACT_BASE_FINAL_NAMES = frozenset({"ABC", "ABCMeta"})


class SourceASTParser:
    """Parse Python source files into ASTCodebase without external dependencies."""

    def parse_directory(
        self,
        path: Path | str,
        name: str = "",
        exclude: list[str] | None = None,
    ) -> ASTCodebase:
        root = Path(path).resolve()
        exclude = exclude or ["venv", ".venv", "__pycache__", ".git", "node_modules"]
        modules = []

        for py_file in root.rglob("*.py"):
            if any(part in py_file.parts for part in exclude):
                continue
            try:
                module = self.parse_file(py_file, root_path=root)
                modules.append(module)
            except SyntaxError, FileNotFoundError:
                # Syntax errors are skipped; FileNotFoundError covers files
                # deleted between rglob enumeration and read (concurrent
                # workers tearing down scratch dirs mid-scan).
                continue

        return ASTCodebase(
            name=name or root.name,
            root_path=str(root),
            modules=modules,
            total_classes=sum(len(m.classes) for m in modules),
            total_functions=sum(len(m.functions) for m in modules),
            total_methods=sum(len(c.methods) for m in modules for c in m.classes),
            total_loc=sum(m.line_count for m in modules),
        )

    def parse_file(
        self, path: Path | str, root_path: Path | str | None = None
    ) -> ModuleInfo:
        file_path = Path(path).resolve()
        source = file_path.read_text(encoding="utf-8")
        tree = ast.parse(source, str(file_path))

        # Calculate line count
        lines = source.splitlines()
        line_count = len(lines)

        # Determine module path from file path
        if root_path:
            root = Path(root_path).resolve()
            try:
                rel_path = file_path.relative_to(root)
                module_path = (
                    str(rel_path.with_suffix("")).replace("/", ".").replace("\\", ".")
                )
            except ValueError:
                module_path = file_path.stem
        else:
            module_path = file_path.stem

        classes = []
        functions = []
        imports = []
        module_vars = []
        docstring = ast.get_docstring(tree)
        all_exports: list[str] = []

        for node in tree.body:
            if isinstance(node, ast.ClassDef):
                classes.append(self._parse_class(node, module_path, str(file_path)))
            elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                functions.append(
                    self._parse_function(node, module_path, str(file_path))
                )
            elif isinstance(node, (ast.Import, ast.ImportFrom)):
                imports.append(self._parse_import(node))
            elif isinstance(node, ast.Assign):
                module_vars.append(self._parse_module_var(node))
            elif (
                isinstance(node, ast.Expr)
                and isinstance(node.value, ast.Constant)
                and isinstance(node.value.value, str)
            ):
                # Module-level string literal (potential docstring)
                pass

        return ModuleInfo(
            path=module_path,
            file_path=str(file_path),
            line_count=line_count,
            classes=classes,
            functions=functions,
            imports=imports,
            module_vars=module_vars,
            docstring=docstring,
            all_exports=all_exports,
        )

    def _parse_class(
        self, node: ast.ClassDef, module_path: str, file_path: str
    ) -> ClassInfo:
        bases = [self._expr_to_str(base) for base in node.bases]

        methods = []
        class_vars = []
        inner_classes = []

        for item in node.body:
            if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                methods.append(
                    self._parse_method(item, module_path, file_path, node.name)
                )
            elif isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
                class_vars.append(
                    ClassVarInfo(
                        name=item.target.id,
                        type_annotation=(
                            self._expr_to_str(item.annotation)
                            if item.annotation
                            else None
                        ),
                        default=self._expr_to_str(item.value) if item.value else None,
                    )
                )
            elif isinstance(item, ast.Assign):
                for target in item.targets:
                    if isinstance(target, ast.Name):
                        class_vars.append(
                            ClassVarInfo(
                                name=target.id,
                                default=(
                                    self._expr_to_str(item.value)
                                    if item.value
                                    else None
                                ),
                            )
                        )
            elif isinstance(item, ast.ClassDef):
                inner_classes.append(self._parse_class(item, module_path, file_path))

        return ClassInfo(
            name=node.name,
            qualified_name=f"{module_path}.{node.name}",
            module_path=module_path,
            file_path=file_path,
            line=node.lineno,
            bases=bases,
            decorators=[self._expr_to_str(d) for d in node.decorator_list],
            methods=methods,
            class_vars=class_vars,
            docstring=ast.get_docstring(node),
            # A class is abstract if it derives from ABC/ABCMeta, or if any of
            # its own methods carry @abstractmethod (checking the class's own
            # decorator_list here would be wrong: @abstractmethod decorates
            # methods, never the class itself).
            is_abstract=(
                any(
                    self._base_final_name(b) in _ABSTRACT_BASE_FINAL_NAMES
                    for b in bases
                )
                or any(m.is_abstract for m in methods)
            ),
            is_dataclass=any(
                self._expr_to_str(d).startswith("dataclass")
                for d in node.decorator_list
            ),
            is_pydantic=any(self._base_final_name(b) == "BaseModel" for b in bases),
            is_protocol=any(b == "Protocol" for b in bases),
            is_enum=any(self._is_enum_base(b) for b in bases),
            inner_classes=inner_classes,
        )

    @staticmethod
    def _base_final_name(base: str) -> str:
        """Return the final dotted segment of a base-class expression.

        Strips any generic subscript first, e.g. "Generic[T]" -> "Generic",
        then returns the segment after the last "." so both bare names
        ("Enum") and dotted/imported names ("enum.Enum") resolve the same way.
        """
        return base.split("[", 1)[0].rsplit(".", 1)[-1]

    @classmethod
    def _is_enum_base(cls, base: str) -> bool:
        """True if a base-class expression marks its class as an Enum.

        Covers the stdlib enum bases (Enum, IntEnum, StrEnum, IntFlag, Flag,
        ReprEnum — bare or dotted, e.g. "enum.Enum") as well as custom Enum/
        Flag subclasses whose name ends in "Enum" or "Flag".
        """
        final = cls._base_final_name(base)
        return (
            final in _ENUM_BASE_FINAL_NAMES
            or final.endswith("Enum")
            or final.endswith("Flag")
        )

    def _parse_method(
        self,
        node: ast.FunctionDef | ast.AsyncFunctionDef,
        module_path: str,
        file_path: str,
        class_name: str,
    ) -> MethodInfo:
        is_property = any(
            self._expr_to_str(d) == "property" for d in node.decorator_list
        )
        is_classmethod = any(
            self._expr_to_str(d) == "classmethod" for d in node.decorator_list
        )
        is_staticmethod = any(
            self._expr_to_str(d) == "staticmethod" for d in node.decorator_list
        )

        return MethodInfo(
            name=node.name,
            qualified_name=f"{module_path}.{class_name}.{node.name}",
            module_path=module_path,
            file_path=file_path,
            line=node.lineno,
            is_async=isinstance(node, ast.AsyncFunctionDef),
            is_method=True,
            is_classmethod=is_classmethod,
            is_staticmethod=is_staticmethod,
            is_property=is_property,
            is_abstract=any(
                self._expr_to_str(d) == "abstractmethod" for d in node.decorator_list
            ),
            decorators=[self._expr_to_str(d) for d in node.decorator_list],
            parameters=self._parse_params(node.args),
            return_type=self._expr_to_str(node.returns) if node.returns else None,
            docstring=ast.get_docstring(node),
            class_name=class_name,
            calls=self._extract_calls(node),
        )

    def _parse_function(
        self,
        node: ast.FunctionDef | ast.AsyncFunctionDef,
        module_path: str,
        file_path: str,
    ) -> FunctionInfo:
        return FunctionInfo(
            name=node.name,
            qualified_name=f"{module_path}.{node.name}",
            module_path=module_path,
            file_path=file_path,
            line=node.lineno,
            is_async=isinstance(node, ast.AsyncFunctionDef),
            decorators=[self._expr_to_str(d) for d in node.decorator_list],
            parameters=self._parse_params(node.args),
            return_type=self._expr_to_str(node.returns) if node.returns else None,
            docstring=ast.get_docstring(node),
        )

    def _parse_params(self, args: ast.arguments) -> list[ParamInfo]:
        params = []
        # Positional args (posonly + regular, in source order). args.defaults
        # holds one entry per *trailing* parameter that has a default — i.e.
        # it right-aligns to the tail of the combined posonly+regular list —
        # so leading parameters without a default get `default=None`.
        positional = args.posonlyargs + args.args
        num_no_default = len(positional) - len(args.defaults)
        for i, arg in enumerate(positional):
            default_node = (
                args.defaults[i - num_no_default] if i >= num_no_default else None
            )
            params.append(
                ParamInfo(
                    name=arg.arg,
                    type_annotation=(
                        self._expr_to_str(arg.annotation) if arg.annotation else None
                    ),
                    default=(
                        self._expr_to_str(default_node)
                        if default_node is not None
                        else None
                    ),
                )
            )
        # Vararg
        if args.vararg:
            params.append(
                ParamInfo(
                    name=args.vararg.arg,
                    type_annotation=(
                        self._expr_to_str(args.vararg.annotation)
                        if args.vararg.annotation
                        else None
                    ),
                    kind="var_positional",
                )
            )
        # Keyword-only args — args.kw_defaults is parallel to args.kwonlyargs
        # (one entry per parameter; None means that specific param has no
        # default), unlike args.defaults which is tail-aligned.
        params.extend(
            ParamInfo(
                name=arg.arg,
                type_annotation=(
                    self._expr_to_str(arg.annotation) if arg.annotation else None
                ),
                default=(
                    self._expr_to_str(default_node)
                    if default_node is not None
                    else None
                ),
                kind="keyword",
            )
            for arg, default_node in zip(args.kwonlyargs, args.kw_defaults, strict=True)
        )
        # Kwarg
        if args.kwarg:
            params.append(
                ParamInfo(
                    name=args.kwarg.arg,
                    type_annotation=(
                        self._expr_to_str(args.kwarg.annotation)
                        if args.kwarg.annotation
                        else None
                    ),
                    kind="var_keyword",
                )
            )
        return params

    def _parse_import(self, node: ast.Import | ast.ImportFrom) -> ImportInfo:
        if isinstance(node, ast.Import):
            names = []
            alias = None
            for alias_node in node.names:
                names.append(alias_node.name)
                if alias_node.asname:
                    alias = alias_node.asname
            return ImportInfo(
                module=names[0] if names else "",
                names=names,
                alias=alias,
                is_from=False,
            )
        else:  # ImportFrom
            names = [alias.name for alias in node.names]
            alias = (
                node.names[0].asname if node.names and node.names[0].asname else None
            )
            return ImportInfo(
                module=node.module or "",
                names=names,
                alias=alias,
                is_from=True,
                is_relative=node.level > 0,
                level=node.level,
            )

    def _parse_module_var(self, node: ast.Assign) -> dict[str, Any]:
        targets = []
        for target in node.targets:
            if isinstance(target, ast.Name):
                targets.append(target.id)
        return {
            "names": targets,
            "value": self._expr_to_str(node.value) if node.value else None,
        }

    def _extract_calls(self, node: ast.AST) -> list[str]:
        calls = []
        for child in ast.walk(node):
            if isinstance(child, ast.Call):
                if isinstance(child.func, ast.Name):
                    calls.append(child.func.id)
                elif isinstance(child.func, ast.Attribute):
                    calls.append(child.func.attr)
        return calls

    def _expr_to_str(self, node: ast.AST | None) -> str:
        if node is None:
            return ""
        if isinstance(node, ast.Name):
            return node.id
        if isinstance(node, ast.Constant):
            return repr(node.value)
        if isinstance(node, ast.Attribute):
            return f"{self._expr_to_str(node.value)}.{node.attr}"
        if isinstance(node, ast.Subscript):
            return f"{self._expr_to_str(node.value)}[{self._expr_to_str(node.slice)}]"
        if isinstance(node, ast.List):
            return f"[{', '.join(self._expr_to_str(e) for e in node.elts)}]"
        if isinstance(node, ast.Tuple):
            return f"({', '.join(self._expr_to_str(e) for e in node.elts)})"
        if isinstance(node, ast.Dict):
            return "{...}"
        if isinstance(node, ast.BinOp):
            return f"{self._expr_to_str(node.left)} ... {self._expr_to_str(node.right)}"
        if isinstance(node, ast.Call):
            func_str = self._expr_to_str(node.func)
            args_str = ", ".join(self._expr_to_str(a) for a in node.args)
            return f"{func_str}({args_str})"
        if sys.version_info >= (3, 9) and isinstance(node, ast.Subscript):
            return f"{self._expr_to_str(node.value)}[{self._expr_to_str(node.slice)}]"
        return ast.unparse(node) if hasattr(ast, "unparse") else ""