orchestrata_tools.diagram.class_diagram

Reference › orchestrata_tools › diagram › class_diagram

This module generates Mermaid class diagrams from an ASTIndex. MermaidDiagram is a small dataclass holding a title, the diagram body and node and edge counts; its render method prefixes the body with a front-matter title block. It is shared with module_diagram.

ClassDiagramGenerator writes each class as a Mermaid class block listing class variables and methods. Method lines carry visibility (- for names beginning with an underscore, + for public and for dunder names), parameter names with type annotations, and the return type. Inheritance edges are drawn from each base, skipping object, BaseModel, Enum and ABC. generate_for_class draws one class with all its subclasses, generate_for_module draws every class of a module, and generate_full draws up to max_classes (default 50) ranked by method count plus subclass count. An unknown class or an empty module gives an empty diagram. The reverse direction is diagram_to_code.

At a glance

Class diagram

Dependency diagram

API

Classes

MermaidDiagram (dataclass)

Decorators: dataclass · Line 14

A generated Mermaid diagram.

ClassDiagramGenerator

Line 28

Generate Mermaid class diagrams from AST classes.

Source

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

"""Generate Mermaid class diagrams from AST."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from orchestrata_tools.models.ast_models import ClassInfo, MethodInfo
    from orchestrata_tools.parser.indexer import ASTIndex


@dataclass
class MermaidDiagram:
    """A generated Mermaid diagram."""

    title: str
    body: str
    node_count: int = 0
    edge_count: int = 0

    def render(self) -> str:
        lines = ["---", f"title: {self.title}", "---", ""]
        lines.append(self.body)
        return "\n".join(lines)


class ClassDiagramGenerator:
    """Generate Mermaid class diagrams from AST classes."""

    def __init__(self, index: ASTIndex):
        self.index = index

    def _escape(self, name: str) -> str:
        """Escape Mermaid special characters."""
        return name.replace("<", "&lt;").replace(">", "&gt;")

    def _visibility(self, method: MethodInfo) -> str:
        if method.name.startswith("__") and method.name.endswith("__"):
            return "+"  # dunder methods are public API
        if method.name.startswith("_"):
            return "-"
        return "+"

    def _method_signature(self, method: MethodInfo) -> str:
        params = ", ".join(
            p.name + (f": {p.type_annotation}" if p.type_annotation else "")
            for p in method.parameters
        )
        ret = f" -> {method.return_type}" if method.return_type else ""
        return f"{self._visibility(method)}{method.name}({params}){ret}"

    def _class_node(self, cls: ClassInfo) -> str:
        lines = [f"    class {cls.name}{{"]
        for cv in cls.class_vars:
            lines.append(f"        {cv.name}")
        for method in cls.methods:
            lines.append(f"        {self._method_signature(method)}")
        lines.append("    }")
        return "\n".join(lines)

    def _inheritance_edges(self, cls: ClassInfo) -> list[str]:
        edges = []
        for base in cls.bases:
            if base in ("object", "BaseModel", "Enum", "ABC"):
                continue
            edges.append(f"    {base} <|-- {cls.name}")
        return edges

    def generate_for_class(self, class_name: str) -> MermaidDiagram:
        """Generate a diagram for a single class and its relationships."""
        cls = self.index.find_class(class_name)
        if not cls:
            return MermaidDiagram(title=f"Class: {class_name}", body="classDiagram\n")

        lines = ["classDiagram"]
        lines.append(self._class_node(cls))
        lines.extend(self._inheritance_edges(cls))

        # Include subclasses
        for sub in self.index.find_subclasses(class_name):
            lines.append(self._class_node(sub))
            lines.extend(self._inheritance_edges(sub))

        body = "\n".join(lines)
        return MermaidDiagram(
            title=f"Class: {class_name}",
            body=body,
            node_count=1 + len(self.index.find_subclasses(class_name)),
            edge_count=len(lines) - 1,
        )

    def generate_for_module(self, module_path: str) -> MermaidDiagram:
        """Generate a diagram for all classes in a module."""
        classes = self.index.find_classes_in_module(module_path)
        if not classes:
            return MermaidDiagram(title=f"Module: {module_path}", body="classDiagram\n")

        lines = ["classDiagram"]
        edges = []
        for cls in classes:
            lines.append(self._class_node(cls))
            edges.extend(self._inheritance_edges(cls))

        lines.extend(edges)
        body = "\n".join(lines)
        return MermaidDiagram(
            title=f"Module: {module_path}",
            body=body,
            node_count=len(classes),
            edge_count=len(edges),
        )

    def generate_full(self, max_classes: int = 50) -> MermaidDiagram:
        """Generate a full class diagram (limited to top classes)."""
        all_classes = self.index.codebase.all_classes
        # Prioritize classes with most methods and most subclasses
        scored = [
            (
                cls,
                cls.method_count + len(self.index.find_subclasses(cls.name)),
            )
            for cls in all_classes
        ]
        scored.sort(key=lambda x: x[1], reverse=True)
        top_classes = [cls for cls, _ in scored[:max_classes]]

        lines = ["classDiagram"]
        edges = []
        for cls in top_classes:
            lines.append(self._class_node(cls))
            edges.extend(self._inheritance_edges(cls))

        lines.extend(edges)
        body = "\n".join(lines)
        return MermaidDiagram(
            title="Full Class Diagram",
            body=body,
            node_count=len(top_classes),
            edge_count=len(edges),
        )