ml_switcheroo.core.tikz.analyser

Static Graph Extractor for TikZ Visualization.

This module implements the analysis logic to convert Python source code (AST) into a language-agnostic Logical Graph. It parses __init__ methods to identify layer definitions (Nodes) and forward/__call__ methods to identify variable data flow (Edges).

The output LogicalGraph is an intermediate representation used by the TikZ emitter to generate visual diagrams.

Classes

GraphExtractor

LibCST Visitor that extracts a LogicalGraph from Python source code.

Module Contents

class ml_switcheroo.core.tikz.analyser.GraphExtractor

Bases: libcst.CSTVisitor

LibCST Visitor that extracts a LogicalGraph from Python source code.

Two-Pass Logic:

1. Init Pass: Scans __init__ or setup to register named layers assigned to self. Populates the node registry.

2. Forward Pass: Scans forward or __call__ to trace variable usage. Builds edges between registered nodes based on data flow.

graph

layer_registry: Dict[str, ml_switcheroo.core.graph.LogicalNode]

provenance: Dict[str, str]

model_name: str = 'GeneratedNet'

visit_ClassDef(node: libcst.ClassDef) → bool | None

Capture the model class name.

leave_Module(original_node: libcst.Module) → None

Finalize graph construction after visiting the whole module. This ensures nodes are populated even if there is no forward pass.

visit_FunctionDef(node: libcst.FunctionDef) → bool | None

Detects entry into lifecycle methods (__init__, forward, etc).

leave_FunctionDef(node: libcst.FunctionDef) → None

Resets context flags upon exiting methods.

visit_Assign(node: libcst.Assign) → bool | None

Handles assignment logic for both layer definition and data flow.

visit_Return(node: libcst.Return) → bool | None

Handles return statements in forward pass to create Output nodes. Also handles case where return contains a functional call directly.