ml_switcheroo.utils.visualizer

AST Visualization Utility.

This module provides the MermaidGenerator, a LibCST visitor that traverses an Abstract Syntax Tree and converts it into a Mermaid.js Graph diagram. It is used to generate “Before” and “After” snapshots of the code structure during transpilation for debugging and documentation.

The generator applies specific branding colors and handles complex node representation (like dotted imports) robustly.

Classes

MermaidGenerator

Generates a Mermaid Graph TD string from a CST Node tree.

Module Contents

class ml_switcheroo.utils.visualizer.MermaidGenerator

Bases: libcst.CSTVisitor

Generates a Mermaid Graph TD string from a CST Node tree.

It traverses the tree and emits nodes and edges formatted with specific branding colors.

COLORS

STYLES = Multiline-String

Show Value

"""
    %% Styles
    classDef default font-family:'Google Sans Normal',color:#20344b,stroke:#20344b;
    classDef modNode fill:#20344b,stroke:#20344b,color:#ffffff,rx:5px,font-family:'Google Sans Medium';
    classDef classNode fill:#ea4335,stroke:#20344b,color:#ffffff,rx:5px,font-family:'Google Sans Medium';
    classDef funcNode fill:#4285f4,stroke:#20344b,color:#ffffff,rx:5px,font-family:'Google Sans Medium';
    classDef callNode fill:#34a853,stroke:#20344b,stroke-width:2px,color:#ffffff,rx:5px,font-family:'Roboto Mono Normal';
    classDef stmtNode fill:#ffffff,stroke:#20344b,stroke-dasharray: 2 2,color:#20344b,font-family:'Roboto Mono Normal';
    classDef argNode fill:#f9ab00,stroke:#20344b,color:#20344b,rx:2px,font-size:10px;
    classDef impNode fill:#57caff,stroke:#20344b,color:#20344b,rx:5px;
    classDef valNode fill:#ff7daf,stroke:#20344b,color:#20344b,rx:5px;
    """

nodes: List[str] = []

edges: List[str] = []

stack: List[str] = []

generate(tree: libcst.CSTNode) → str

Converts a CST Node into a Mermaid graph definition string.

Parameters:

tree (cst.CSTNode) – The root node of the tree to visualize.

Returns:

A complete Mermaid.js graph definition including styles.

Return type:

str

visit_Module(node: libcst.Module) → bool | None

Visits Module root.

leave_Module(node: libcst.Module) → None

Leaves Module root.

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

Visits Class Definitions.

leave_ClassDef(node: libcst.ClassDef) → None

Leaves Class Definitions.

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

Visits Function Definitions.

leave_FunctionDef(node: libcst.FunctionDef) → None

Leaves Function Definitions.

visit_Call(node: libcst.Call) → bool | None

Visits Function Calls.

leave_Call(node: libcst.Call) → None

Leaves Function Calls.

visit_Arg(node: libcst.Arg) → bool | None

Visits Arguments inside a Call.

leave_Arg(node: libcst.Arg) → None

Leaves Arguments.

visit_Import(node: libcst.Import) → bool | None

Visits Import statements (collapsing them into single nodes).

visit_ImportFrom(node: libcst.ImportFrom) → bool | None

Visits From-Import statements (collapsed).

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

Visits Assignment statements.

leave_Assign(node: libcst.Assign) → None

Leaves Assignment statements.

visit_SimpleString(node: libcst.SimpleString) → bool | None

Visits Strings to allow visualizing constant values directly in the graph.