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::

Text Only

>>> 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')",
    )

TraceKwargs

Bases: TypedDict

Keyword arguments forwarded to :func:trace_flow.

Public mirror of the kwargs accepted by :func:trace_flow so call sites (notably :mod:axm_ast.hooks.flows) can build typed kwargs dicts without redefining the contract locally.

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

class TraceKwargs(TypedDict):
    """Keyword arguments forwarded to :func:`trace_flow`.

    Public mirror of the kwargs accepted by :func:`trace_flow` so call sites
    (notably :mod:`axm_ast.hooks.flows`) can build typed kwargs dicts without
    redefining the contract locally.
    """

    max_depth: int
    cross_module: bool
    detail: str
    exclude_stdlib: bool

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[Tree, 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[Tree, 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_callees_workspace(ws, symbol)

Find all callees of a symbol across a workspace.

Searches every package in the workspace for callees of the given symbol. Module names are prefixed with pkg_name:: for disambiguation.

Parameters:

Name	Type	Description	Default
ws	WorkspaceInfo	Analyzed workspace info.	required
symbol	str	Name of the function/method to inspect.	required

Returns:

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

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

def find_callees_workspace(
    ws: WorkspaceInfo,
    symbol: str,
) -> list[CallSite]:
    """Find all callees of a symbol across a workspace.

    Searches every package in the workspace for callees of the
    given symbol. Module names are prefixed with ``pkg_name::``
    for disambiguation.

    Args:
        ws: Analyzed workspace info.
        symbol: Name of the function/method to inspect.

    Returns:
        List of CallSite objects for each call made by the symbol.
    """
    all_callees: list[CallSite] = []
    cache: dict[str, tuple[Tree, str]] = {}

    for pkg in ws.packages:
        callees = find_callees(pkg, symbol, _parse_cache=cache)
        for call in callees:
            call.module = f"{pkg.name}::{call.module}"
            all_callees.append(call)

    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_flow_compact(steps)

Format flow steps as a compact tree with box-drawing characters.

Each step is rendered on one line. Depth-0 is the root (no prefix), deeper levels use box-drawing connectors with indentation proportional to depth.

Parameters:

Name	Type	Description	Default
steps	list[FlowStep]	Ordered list of FlowSteps (BFS order, ascending depth).	required

Returns:

Type	Description
str	Tree-formatted string. Empty string when steps is empty.

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

def format_flow_compact(steps: list[FlowStep]) -> str:
    """Format flow steps as a compact tree with box-drawing characters.

    Each step is rendered on one line.  Depth-0 is the root (no prefix),
    deeper levels use box-drawing connectors with
    indentation proportional to depth.

    Args:
        steps: Ordered list of FlowSteps (BFS order, ascending depth).

    Returns:
        Tree-formatted string.  Empty string when *steps* is empty.
    """
    if not steps:
        return ""
    return "\n".join(_format_step_line(steps, i, s) for i, s in enumerate(steps))

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; "compact" produces a tree-formatted string. Must be one of VALID_DETAILS; raises :exc:ValueError otherwise.	'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]	Tuple of (steps, truncated) where steps is a list of FlowStep
bool	objects ordered by depth then discovery, and truncated is True
tuple[list[FlowStep], bool]	when at least one frontier node at max_depth had unexpanded
tuple[list[FlowStep], bool]	children.

Example

steps, truncated = 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,
) -> tuple[list[FlowStep], bool]:
    """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; ``"compact"``
            produces a tree-formatted string.  Must be one of
            ``VALID_DETAILS``; raises :exc:`ValueError` otherwise.
        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:
        Tuple of (steps, truncated) where *steps* is a list of FlowStep
        objects ordered by depth then discovery, and *truncated* is True
        when at least one frontier node at *max_depth* had unexpanded
        children.

    Example:
        >>> steps, truncated = trace_flow(pkg, "main", max_depth=3)
        >>> for s in steps:
        ...     print(f"{'  ' * s.depth}{s.name} ({s.module}:{s.line})")
    """
    if detail not in VALID_DETAILS:
        msg = f"Invalid detail={detail!r}; must be one of {sorted(VALID_DETAILS)}"
        raise ValueError(msg)

    t0 = time.perf_counter()

    # Find the entry point location
    entry_mod, entry_line = _find_symbol_location(pkg, entry)
    if entry_mod is None:
        msg = f"Symbol {entry!r} not found in package"
        raise ValueError(msg)

    # 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,
        exclude_stdlib=exclude_stdlib,
        pkg_symbols=pkg_symbols,
    )

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

    truncated = False

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

        if depth >= max_depth:
            truncated = truncated or _check_frontier_truncated(
                current_mod,
                current,
                current_pkg,
                callee_index,
                ctx,
                exclude_stdlib=exclude_stdlib,
                pkg_symbols=pkg_symbols,
                visited=visited,
            )
            continue

        callees = _get_callees(current_mod, current, current_pkg, callee_index, ctx)
        _process_local_callees(
            callees=callees,
            exclude_stdlib=exclude_stdlib,
            pkg_symbols=pkg_symbols,
            visited=visited,
            current_chain=current_chain,
            depth=depth,
            steps=steps,
            queue=queue,
            current_pkg=current_pkg,
        )

        if cross_module:
            _resolve_cross_module_callees(
                callees,
                _ResolutionScope(
                    current_mod=current_mod,
                    current_pkg=current_pkg,
                    original_pkg=pkg,
                    depth=depth,
                    current_chain=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, truncated