Flows

flows

Execution flow tracing via entry point detection and BFS call graph.

Detects framework-specific entry points (cyclopts, click, Flask, FastAPI, pytest, __main__ guards) and traces execution flows through the call graph using BFS.

Example::

>>> from axm_ast.core.analyzer import analyze_package
>>> from axm_ast.core.flows import find_entry_points, trace_flow
>>> pkg = analyze_package(Path("src/mylib"))
>>> entries = find_entry_points(pkg)
>>> for e in entries:
...     print(f"{e.framework}: {e.name} ({e.module}:{e.line})")

EntryPoint

Bases: BaseModel

A detected entry point in the codebase.

Source code in packages/axm-ast/src/axm_ast/core/flows.py

class EntryPoint(BaseModel):
    """A detected entry point in the codebase."""

    model_config = ConfigDict(extra="forbid")

    name: str
    module: str
    kind: str  # "decorator", "test", "main_guard", "export"
    line: int
    framework: str  # "cyclopts", "click", "flask", "fastapi", "pytest", "main", "all"

FlowStep

Bases: BaseModel

A single step in a traced execution flow.

Source code in packages/axm-ast/src/axm_ast/core/flows.py

class FlowStep(BaseModel):
    """A single step in a traced execution flow."""

    model_config = ConfigDict(extra="forbid")

    name: str
    module: str
    line: int
    depth: int
    chain: list[str]
    resolved_module: str | None = Field(
        default=None,
        description="Dotted module path when resolved across modules",
    )
    source: str | None = Field(
        default=None,
        description="Source code of the symbol (when detail='source')",
    )

build_callee_index(pkg)

Pre-compute a callee index for the entire package in one pass.

Instead of scanning all modules per symbol (O(modules x AST) per BFS step), this builds a {(module, symbol): [CallSite]} dict in a single pass. BFS then uses O(1) dict lookups.

Parameters:

Name	Type	Description	Default
pkg	PackageInfo	Analyzed package info.	required

Returns:

Type	Description
dict[tuple[str, str], list[CallSite]]	Dict mapping (module_dotted_name, function_name) to callees.

Source code in packages/axm-ast/src/axm_ast/core/flows.py

def build_callee_index(
    pkg: PackageInfo,
) -> dict[tuple[str, str], list[CallSite]]:
    """Pre-compute a callee index for the entire package in one pass.

    Instead of scanning all modules per symbol (O(modules x AST) per BFS step),
    this builds a ``{(module, symbol): [CallSite]}`` dict in a single pass.
    BFS then uses O(1) dict lookups.

    Args:
        pkg: Analyzed package info.

    Returns:
        Dict mapping ``(module_dotted_name, function_name)`` to callees.
    """
    index: dict[tuple[str, str], list[CallSite]] = {}

    for mod in pkg.modules:
        mod_name = module_dotted_name(mod.path, pkg.root)
        source = mod.path.read_text(encoding="utf-8")
        tree = parse_source(source)

        func_ranges = _find_all_function_ranges(tree.root_node)
        for fr in func_ranges:
            calls = _extract_scoped_calls(tree.root_node, mod_name, source, fr)
            index[(mod_name, fr.name)] = calls

    return index

find_callees(pkg, symbol, *, _parse_cache=None)

Find all functions called by a given symbol (forward call graph).

This is the inverse of find_callers: instead of asking "who calls X?", it asks "what does X call?".

Parameters:

Name	Type	Description	Default
pkg	PackageInfo	Analyzed package info.	required
symbol	str	Name of the function/method to inspect.	required
_parse_cache	dict[str, tuple[Any, str]] | None	Optional cache of {path_str: (tree, source)} to avoid re-parsing the same file. Shared across BFS iterations in trace_flow.	None

Returns:

Type	Description
list[CallSite]	List of CallSite objects for each call made by the symbol.

Example

callees = find_callees(pkg, "main") for c in callees: ... print(f" calls {c.symbol} at {c.module}:{c.line}")

Source code in packages/axm-ast/src/axm_ast/core/flows.py

def find_callees(
    pkg: PackageInfo,
    symbol: str,
    *,
    _parse_cache: dict[str, tuple[Any, str]] | None = None,
) -> list[CallSite]:
    """Find all functions called by a given symbol (forward call graph).

    This is the inverse of ``find_callers``: instead of asking "who calls X?",
    it asks "what does X call?".

    Args:
        pkg: Analyzed package info.
        symbol: Name of the function/method to inspect.
        _parse_cache: Optional cache of ``{path_str: (tree, source)}``
            to avoid re-parsing the same file.  Shared across BFS
            iterations in ``trace_flow``.

    Returns:
        List of CallSite objects for each call made by the symbol.

    Example:
        >>> callees = find_callees(pkg, "main")
        >>> for c in callees:
        ...     print(f"  calls {c.symbol} at {c.module}:{c.line}")
    """
    cache = _parse_cache if _parse_cache is not None else {}
    all_callees: list[CallSite] = []

    for mod in pkg.modules:
        mod_name = module_dotted_name(mod.path, pkg.root)
        path_key = str(mod.path)

        if path_key in cache:
            tree, source = cache[path_key]
        else:
            source = mod.path.read_text(encoding="utf-8")
            tree = parse_source(source)
            cache[path_key] = (tree, source)

        # Find the function definition node for this symbol
        func_ranges = _find_function_nodes(tree.root_node, symbol)
        for func_range in func_ranges:
            # Scope call extraction to this function's subtree
            calls = _extract_scoped_calls(tree.root_node, mod_name, source, func_range)
            all_callees.extend(calls)

    return all_callees

find_entry_points(pkg)

Detect framework-registered entry points across a package.

Scans for:

• Decorator-based: cyclopts, click, Flask, FastAPI
• Test functions: test_* prefix
• Main guards: if __name__ == "__main__" blocks
• __all__ exports

Parameters:

Name	Type	Description	Default
pkg	PackageInfo	Analyzed package info.	required

Returns:

Type	Description
list[EntryPoint]	List of EntryPoint objects sorted by module then line.

Source code in packages/axm-ast/src/axm_ast/core/flows.py

def find_entry_points(pkg: PackageInfo) -> list[EntryPoint]:
    """Detect framework-registered entry points across a package.

    Scans for:

    - **Decorator-based**: cyclopts, click, Flask, FastAPI
    - **Test functions**: ``test_*`` prefix
    - **Main guards**: ``if __name__ == "__main__"`` blocks
    - **``__all__`` exports**

    Args:
        pkg: Analyzed package info.

    Returns:
        List of EntryPoint objects sorted by module then line.
    """
    entries: list[EntryPoint] = []

    for mod in pkg.modules:
        mod_name = module_dotted_name(mod.path, pkg.root)
        entries.extend(_scan_module_entries(mod, mod_name))

    entries.sort(key=lambda e: (e.module, e.line))
    return entries

format_flows(entry_points)

Format entry point results as human-readable grouped output.

Parameters:

Name	Type	Description	Default
entry_points	list[EntryPoint]	List of detected entry points.	required

Returns:

Type	Description
str	Formatted string grouped by framework.

Source code in packages/axm-ast/src/axm_ast/core/flows.py

def format_flows(entry_points: list[EntryPoint]) -> str:
    """Format entry point results as human-readable grouped output.

    Args:
        entry_points: List of detected entry points.

    Returns:
        Formatted string grouped by framework.
    """
    if not entry_points:
        return "✅ No entry points detected."

    # Group by framework
    by_framework: dict[str, list[EntryPoint]] = {}
    for ep in entry_points:
        by_framework.setdefault(ep.framework, []).append(ep)

    lines: list[str] = [f"🔍 {len(entry_points)} entry point(s) detected:\n"]

    for framework, eps in sorted(by_framework.items()):
        lines.append(f"  📦 {framework} ({len(eps)}):")
        for ep in eps:
            lines.append(f"    • {ep.name} ({ep.module}:{ep.line}) [{ep.kind}]")
        lines.append("")

    return "\n".join(lines)

trace_flow(pkg, entry, *, max_depth=5, cross_module=False, detail='trace', callee_index=None, exclude_stdlib=True)

Trace execution flow from an entry point via BFS.

Follows the forward call graph from entry up to max_depth levels deep. Uses a visited set to handle circular calls.

Parameters:

Name	Type	Description	Default
pkg	PackageInfo	Analyzed package info.	required
entry	str	Name of the entry point function to trace from.	required
max_depth	int	Maximum BFS depth (default 5).	5
cross_module	bool	If True, resolve imports and continue BFS into external modules on-demand.	False
detail	str	Level of detail — "trace" (default) returns names and positions only; "source" enriches each step with the function's source code.	'trace'
callee_index	dict[tuple[str, str], list[CallSite]] | None	Optional pre-computed index from :func:build_callee_index. When provided, BFS uses O(1) dict lookups instead of scanning all modules.	None
exclude_stdlib	bool	If True (default), skip callees whose name matches a stdlib module or Python builtin (e.g. len, isinstance). Set to False to include them.	True

Returns:

Type	Description
list[FlowStep]	List of FlowStep objects ordered by depth then discovery.

Example

steps = trace_flow(pkg, "main", max_depth=3) for s in steps: ... print(f"{' ' * s.depth}{s.name} ({s.module}:{s.line})")

Source code in packages/axm-ast/src/axm_ast/core/flows.py

def trace_flow(  # noqa: PLR0913
    pkg: PackageInfo,
    entry: str,
    *,
    max_depth: int = 5,
    cross_module: bool = False,
    detail: str = "trace",
    callee_index: dict[tuple[str, str], list[CallSite]] | None = None,
    exclude_stdlib: bool = True,
) -> list[FlowStep]:
    """Trace execution flow from an entry point via BFS.

    Follows the forward call graph from *entry* up to *max_depth*
    levels deep. Uses a visited set to handle circular calls.

    Args:
        pkg: Analyzed package info.
        entry: Name of the entry point function to trace from.
        max_depth: Maximum BFS depth (default 5).
        cross_module: If True, resolve imports and continue BFS
            into external modules on-demand.
        detail: Level of detail — ``"trace"`` (default) returns
            names and positions only; ``"source"`` enriches each
            step with the function's source code.
        callee_index: Optional pre-computed index from
            :func:`build_callee_index`.  When provided, BFS uses
            O(1) dict lookups instead of scanning all modules.
        exclude_stdlib: If True (default), skip callees whose name
            matches a stdlib module or Python builtin (e.g. ``len``,
            ``isinstance``).  Set to False to include them.

    Returns:
        List of FlowStep objects ordered by depth then discovery.

    Example:
        >>> steps = trace_flow(pkg, "main", max_depth=3)
        >>> for s in steps:
        ...     print(f"{'  ' * s.depth}{s.name} ({s.module}:{s.line})")
    """
    t0 = time.perf_counter()

    # Find the entry point location
    entry_mod, entry_line = _find_symbol_location(pkg, entry)
    if entry_mod is None:
        return []

    # Pre-compute set of symbols defined in the package so we can
    # distinguish project callees from stdlib method calls (e.g.
    # logger.info → "info" is not in pkg_symbols → skip).
    pkg_symbols = _build_package_symbols(pkg) if exclude_stdlib else frozenset()

    steps: list[FlowStep] = []
    # Use (module, symbol) tuples to handle same-named symbols
    # in different modules.
    visited: set[tuple[str, str]] = {(entry_mod, entry)}
    # Queue: (symbol, depth, chain, source_pkg, source_module_dotted)
    queue: deque[tuple[str, int, list[str], PackageInfo, str]] = deque()
    queue.append((entry, 0, [entry], pkg, entry_mod))

    # Shared BFS context for cross-module resolution
    ctx = _CrossModuleContext(
        visited=visited,
        queue=queue,
        steps=steps,
        detail=detail,
    )

    # Add the entry point itself
    steps.append(
        FlowStep(
            name=entry,
            module=entry_mod,
            line=entry_line,
            depth=0,
            chain=[entry],
        )
    )

    while queue:
        current, depth, current_chain, current_pkg, current_mod = queue.popleft()

        if depth >= max_depth:
            continue

        if callee_index is not None:
            callees = callee_index.get((current_mod, current), [])
        else:
            callees = find_callees(current_pkg, current, _parse_cache=ctx.parse_cache)
        for callee in callees:
            if exclude_stdlib and (
                _is_stdlib_or_builtin(callee.symbol) or callee.symbol not in pkg_symbols
            ):
                continue
            callee_key = (callee.module, callee.symbol)
            if callee_key not in visited:
                visited.add(callee_key)
                new_chain = [*current_chain, callee.symbol]
                steps.append(
                    FlowStep(
                        name=callee.symbol,
                        module=callee.module,
                        line=callee.line,
                        depth=depth + 1,
                        chain=new_chain,
                    )
                )
                queue.append(
                    (callee.symbol, depth + 1, new_chain, current_pkg, callee.module)
                )

        if not cross_module:
            continue

        # Cross-module resolution: find symbols that were imported
        # but not defined locally.
        _resolve_cross_module_callees(
            callees,
            current_mod,
            current_pkg,
            pkg,
            depth,
            current_chain,
            ctx,
        )

    if detail == "source":
        _enrich_steps_with_source(steps, pkg)

    elapsed = time.perf_counter() - t0
    logger.debug(
        "Traced %s in %.2fs (%d steps, depth=%d)",
        entry,
        elapsed,
        len(steps),
        max_depth,
    )

    return steps