orchestrata_tools.diagram.diagram_to_code

Reference › orchestrata_tools › diagram › diagram_to_code

This module is the reverse direction of the diagram pipeline: it reads a Mermaid classDiagram and generates Python code skeletons from it. It is a port of DiagramToCode from the old diagram_generator.py script and is used by the round-trip verifier and the spec compliance audit. It is imported by the package __init__ and diagram; tests/test_diagram_to_code_roundtrip.py and tests/test_orchestrata_tools_diagram_to_code.py cover it. The forward direction is class_diagram.

ParsedClass describes one class read from a diagram: its name, stereotypes, decorators, variables (as ClassVarInfo), methods (as FunctionInfo), bases, implemented interfaces, and relationship lists for compositions, aggregations, dependencies and associations. It reuses the models in ast_models.

DiagramToCode has three methods. parse turns diagram text into a list of ParsedClass. generate_code renders the classes as Python source, either as one module named by module_name or, with per_module=True, as a dictionary of module name to source; imports can be omitted. write_modules writes per-module files into an output directory and returns the file contents keyed by path.

At a glance

Class diagram

Dependency diagram

API

Classes

ParsedClass (dataclass)

Decorators: dataclass · Line 29

A class parsed from a Mermaid classDiagram.

DiagramToCode

Line 49

Parse Mermaid class diagrams and generate Python code skeletons.

Source

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

"""Parse Mermaid class diagrams and generate Python code skeletons.

Reverse direction of the diagram pipeline: turns ``classDiagram`` blocks
into structured ``ParsedClass`` objects, then generates Python modules from them.

Used by the round-trip verifier and spec compliance auditor
(``generate_comprehensive_audit.py``).

Port of ``DiagramToCode`` from ``scripts/diagram_generator.py``.
"""

from __future__ import annotations

import os
import re
from collections import OrderedDict
from dataclasses import dataclass, field

from orchestrata_tools.models.ast_models import (
    ClassVarInfo,
    FunctionInfo,
    ParamInfo,
)

# ── ParsedClass ───────────────────────────────────────────────────────


@dataclass
class ParsedClass:
    """A class parsed from a Mermaid classDiagram."""

    name: str
    stereotypes: list[str] = field(default_factory=list)
    decorators: list[str] = field(default_factory=list)
    variables: list[ClassVarInfo] = field(default_factory=list)
    methods: list[FunctionInfo] = field(default_factory=list)
    bases: list[str] = field(default_factory=list)
    implements: list[str] = field(default_factory=list)
    compositions: list[tuple[str, str]] = field(default_factory=list)
    aggregations: list[tuple[str, str]] = field(default_factory=list)
    dependencies: list[tuple[str, str]] = field(default_factory=list)
    associations: list[tuple[str, str]] = field(default_factory=list)
    namespace: str = ""


# ── DiagramToCode ─────────────────────────────────────────────────────


class DiagramToCode:
    """Parse Mermaid class diagrams and generate Python code skeletons."""

    _BUILTIN_TYPE_NAMES = frozenset(
        {
            "str",
            "int",
            "float",
            "bool",
            "bytes",
            "None",
            "Any",
            "Callable",
            "Type",
            "ClassVar",
            "Optional",
            "List",
            "Dict",
            "Set",
            "Tuple",
            "list",
            "dict",
            "set",
            "tuple",
            "frozenset",
            "datetime",
            "Path",
            "UUID",
            "uuid4",
            "timedelta",
            "object",
            "type",
            "Exception",
        }
    )

    # ── helpers ────────────────────────────────────────────────────

    @staticmethod
    def _class_name_to_filename(name: str) -> str:
        s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", name)
        s = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", s)
        return s.lower()

    @staticmethod
    def _mermaid_type_to_python(type_str: str) -> str:
        if not type_str:
            return type_str
        result = []
        tilde_count = 0
        for char in type_str:
            if char == "~":
                if tilde_count % 2 == 0:
                    result.append("[")
                else:
                    result.append("]")
                tilde_count += 1
            else:
                result.append(char)
        return "".join(result)

    @staticmethod
    def _split_params(params_str: str) -> list[str]:
        parts = []
        depth = 0
        current = []
        for ch in params_str:
            if ch in ("(", "[", "{"):
                depth += 1
                current.append(ch)
            elif ch in (")", "]", "}"):
                depth -= 1
                current.append(ch)
            elif ch == "," and depth == 0:
                parts.append("".join(current))
                current = []
            else:
                current.append(ch)
        if current:
            parts.append("".join(current))
        return parts

    @staticmethod
    def _split_on_first_eq(s: str) -> list[str]:
        depth = 0
        for i, ch in enumerate(s):
            if ch in ("(", "[", "{"):
                depth += 1
            elif ch in (")", "]", "}"):
                depth -= 1
            elif ch == "=" and depth == 0:
                return [s[:i], s[i + 1 :]]
        return [s]

    def _extract_type_references(self, pc: ParsedClass) -> set[str]:
        raw: list[str] = []
        for var in pc.variables:
            if var.type_annotation:
                raw.append(var.type_annotation)
        for method in pc.methods:
            for p in method.parameters:
                if p.type_annotation:
                    raw.append(p.type_annotation)
            if method.return_type:
                raw.append(method.return_type)
        refs: set[str] = set()
        for ann in raw:
            for token in re.findall(r"\b([A-Z][A-Za-z0-9_]*)\b", ann):
                if token not in self._BUILTIN_TYPE_NAMES:
                    refs.add(token)
        refs.discard(pc.name)
        return refs

    def _is_enum_class(self, pc: ParsedClass) -> bool:
        return any("enum" in s.lower() for s in pc.stereotypes + pc.decorators)

    # ── parsing ────────────────────────────────────────────────────

    def parse(self, diagram: str) -> list[ParsedClass]:
        lines = diagram.splitlines()
        classes: dict[str, ParsedClass] = {}
        class_order: list[str] = []
        current_class: str | None = None
        in_class_body = False
        current_namespace: str | None = None

        for line in lines:
            stripped = line.strip()

            if stripped.startswith("classDiagram") or stripped.startswith(
                "classDiagram-direction"
            ):
                continue

            ns_match = re.match(r"namespace\s+([\w.]+)\s*\{", stripped)
            if ns_match:
                current_namespace = ns_match.group(1)
                continue

            class_match = re.match(r"class\s+(\w+)(?:\s*<<(.+?)>>)?\s*\{", stripped)
            if class_match:
                cname = class_match.group(1)
                stereo = class_match.group(2)
                pc = ParsedClass(name=cname, namespace=current_namespace or "")
                if stereo:
                    pc.stereotypes = [s.strip() for s in stereo.split(",")]
                current_class = cname
                in_class_body = True
                classes[cname] = pc
                class_order.append(cname)
                continue

            if stripped == "}" and in_class_body:
                in_class_body = False
                current_class = None
                continue

            if stripped == "}" and current_namespace and not in_class_body:
                current_namespace = None
                continue

            if in_class_body and current_class:
                self._parse_class_member(classes[current_class], stripped)
                continue

            if not in_class_body:
                self._parse_relationship(stripped, classes)

        return [classes[n] for n in class_order if n in classes]

    def _parse_class_member(self, pc: ParsedClass, line: str) -> None:
        if line == "---":
            return

        decorator_match = re.match(r"<<(.+?)>>", line)
        if decorator_match:
            pc.decorators.append(decorator_match.group(1))
            return

        method_match = re.match(
            r"(?:\{(.+?)\}\s*)?([+#-]?\s*)(async\s+)?(\w+)\(",
            line,
        )
        if method_match:
            decos_str = method_match.group(1) or ""
            is_async = bool(method_match.group(3))
            name = method_match.group(4)

            rest = line[method_match.end() :]
            params_str, return_type = self._split_params_and_return(rest)
            if return_type:
                return_type = self._mermaid_type_to_python(return_type)
                # Strip leading "->" if the parser already captured it
                return_type = re.sub(r"^\s*->\s*", "", return_type)

            decorators = (
                [d.strip().lstrip("@") for d in decos_str.split() if d.strip()]
                if decos_str
                else []
            )

            parameters = self._parse_params_sig(params_str)

            is_classmethod = "classmethod" in decorators
            is_staticmethod = "staticmethod" in decorators
            is_property = "property" in decorators
            is_abstract = "abstractmethod" in decorators

            fi = FunctionInfo(
                name=name,
                qualified_name=f"{pc.name}.{name}",
                module_path="",
                file_path="",
                line=0,
                is_async=is_async,
                decorators=decorators,
                parameters=parameters,
                return_type=return_type,
            )
            # Attach extra metadata via a private dict for code generation
            fi._extra = {  # type: ignore[attr-defined]
                "is_method": not is_staticmethod,
                "is_classmethod": is_classmethod,
                "is_staticmethod": is_staticmethod,
                "is_property": is_property,
                "is_abstract": is_abstract,
                "class_name": pc.name,
            }
            pc.methods.append(fi)
            return

        var_match = re.match(
            r"([#$+-]?\s*)(\w+)(?::\s*(.+?))?(?:\s*=\s*(.+))?$",
            line,
        )
        if var_match:
            vname = var_match.group(2)
            annotation = var_match.group(3)
            if annotation:
                annotation = self._mermaid_type_to_python(annotation)
            default = var_match.group(4)

            cvi = ClassVarInfo(
                name=vname,
                type_annotation=annotation,
                default=default,
            )
            pc.variables.append(cvi)

    def _split_params_and_return(self, rest: str) -> tuple[str, str | None]:
        depth = 0
        close_idx = -1
        for i, ch in enumerate(rest):
            if ch == "(":
                depth += 1
            elif ch == ")":
                if depth == 0:
                    close_idx = i
                    break
                depth -= 1

        if close_idx == -1:
            return rest, None

        params_str = rest[:close_idx]
        after = rest[close_idx + 1 :].strip()
        return_type = None
        if after.startswith(":"):
            return_type = after[1:].strip() or None
        elif after:
            return_type = after.strip() or None
        return params_str, return_type

    def _parse_params_sig(self, params_str: str) -> list[ParamInfo]:
        if not params_str.strip():
            return []

        params = []
        for p in self._split_params(params_str):
            p = p.strip()
            if not p:
                continue

            kind = "positional"
            default = None

            if "=" in p:
                parts = self._split_on_first_eq(p)
                p = parts[0].strip()
                default = parts[1].strip()

            if ":" in p:
                parts = p.split(":", 1)
                name = parts[0].strip()
                type_ann = self._mermaid_type_to_python(parts[1].strip())
            else:
                name = p
                type_ann = None

            if name.startswith("*"):
                if name.startswith("**"):
                    name = name[2:]
                    kind = "kwargs"
                else:
                    name = name[1:]
                    kind = "varargs"

            params.append(
                ParamInfo(
                    name=name,
                    type_annotation=type_ann,
                    default=default,
                    kind=kind,
                )
            )

        return params

    def _parse_relationship(self, line: str, classes: dict[str, ParsedClass]) -> None:
        rel_patterns = [
            (r'(\w+)\s+<\|--\s+(\w+)(?:\s+"(.+?)")?', "inheritance"),
            (r'(\w+)\s+\.\.\|>\s+(\w+)(?:\s+"(.+?)")?', "implementation"),
            (r'(\w+)\s+\*--\s+(\w+)(?:\s+"(.+?)")?', "composition"),
            (r'(\w+)\s+o--\s+(\w+)(?:\s+"(.+?)")?', "aggregation"),
            (r'(\w+)\s+\.\.>\s+(\w+)(?:\s+"(.+?)")?', "dependency"),
            (r'(\w+)\s+-->\s+(\w+)(?:\s+"(.+?)")?', "association"),
            (r'(\w+)\s+--\s+(\w+)(?:\s+"(.+?)")?', "dir_association"),
        ]

        for pattern, rel_type in rel_patterns:
            match = re.match(pattern, line)
            if match:
                from_class = match.group(1)
                to_class = match.group(2)
                label = match.group(3) or ""

                if from_class not in classes:
                    classes[from_class] = ParsedClass(name=from_class)
                if to_class not in classes:
                    classes[to_class] = ParsedClass(name=to_class)

                if rel_type == "inheritance":
                    classes[to_class].bases.append(from_class)
                elif rel_type == "implementation":
                    classes[to_class].implements.append(from_class)
                elif rel_type == "composition":
                    classes[from_class].compositions.append((to_class, label))
                elif rel_type == "aggregation":
                    classes[from_class].aggregations.append((to_class, label))
                elif rel_type == "dependency":
                    classes[from_class].dependencies.append((to_class, label))
                elif rel_type == "association":
                    classes[from_class].associations.append((to_class, label))
                break

    # ── code generation ────────────────────────────────────────────

    def generate_code(
        self,
        parsed: list[ParsedClass],
        *,
        module_name: str = "scaffold",
        include_imports: bool = True,
        per_module: bool = False,
    ) -> str | dict[str, str]:
        if per_module:
            return self._generate_code_per_module(
                parsed, include_imports=include_imports
            )

        lines = [
            f'"""Generated from Mermaid diagram — {module_name} module."""',
            "",
        ]
        if include_imports:
            import_lines = self._generate_imports(parsed)
            if import_lines:
                lines.extend(import_lines)
                lines.append("")
        for pc in parsed:
            lines.extend(self._generate_class(pc))
            lines.append("")
        return "\n".join(lines)

    def _generate_code_per_module(
        self, parsed: list[ParsedClass], *, include_imports: bool = True
    ) -> dict[str, str]:
        groups: OrderedDict[str, list[ParsedClass]] = OrderedDict()
        for pc in parsed:
            ns = pc.namespace or ""
            groups.setdefault(ns, []).append(pc)

        result: dict[str, str] = {}
        for ns, group in groups.items():
            module_path = ns.replace(".", "/") if ns else "scaffold"
            lines = [
                f'"""Generated from Mermaid diagram — {ns or "scaffold"} module."""',
                "",
            ]
            if include_imports:
                il = self._generate_imports(group)
                if il:
                    lines.extend(il)
                    lines.append("")
            for pc in group:
                lines.extend(self._generate_class(pc))
                lines.append("")
            result[module_path] = "\n".join(lines)
        return result

    def _generate_imports(self, parsed: list[ParsedClass]) -> list[str]:
        imports = set()
        for pc in parsed:
            if any("enum" in s.lower() for s in pc.stereotypes + pc.decorators):
                imports.add("from enum import Enum")
            if any("dataclass" in s for s in pc.stereotypes + pc.decorators):
                imports.add("from dataclasses import dataclass, field")
            if any("pydantic" in s for s in pc.stereotypes + pc.decorators):
                imports.add("from pydantic import BaseModel")
            if any("abstract" in s.lower() for s in pc.stereotypes + pc.decorators):
                imports.add("from abc import ABC, abstractmethod")
            if pc.implements:
                imports.add("from typing import Protocol")
            for method in pc.methods:
                if method.is_async:
                    imports.add("import asyncio")
                for p in method.parameters:
                    if p.type_annotation:
                        for t in (
                            "Optional",
                            "List",
                            "Dict",
                            "Set",
                            "Tuple",
                            "Callable",
                            "Any",
                        ):
                            if t in (p.type_annotation or ""):
                                imports.add(f"from typing import {t}")
                if method.return_type:
                    for t in (
                        "Optional",
                        "List",
                        "Dict",
                        "Set",
                        "Tuple",
                        "Callable",
                        "Any",
                    ):
                        if t in (method.return_type or ""):
                            imports.add(f"from typing import {t}")
            for var in pc.variables:
                if var.type_annotation and "ClassVar" in var.type_annotation:
                    imports.add("from typing import ClassVar")

        stdlib_type_imports = {
            "datetime": "from datetime import datetime",
            "Path": "from pathlib import Path",
            "UUID": "from uuid import UUID",
            "uuid4": "from uuid import uuid4",
            "timedelta": "from datetime import timedelta",
        }
        all_annotations: list[str] = []
        for pc in parsed:
            for var in pc.variables:
                if var.type_annotation:
                    all_annotations.append(var.type_annotation)
            for method in pc.methods:
                for p in method.parameters:
                    if p.type_annotation:
                        all_annotations.append(p.type_annotation)
                if method.return_type:
                    all_annotations.append(method.return_type)
        combined = " ".join(all_annotations)
        for type_name, import_line in stdlib_type_imports.items():
            if type_name in combined:
                imports.add(import_line)

        return sorted(imports)

    def _generate_class(self, pc: ParsedClass) -> list[str]:
        lines = []
        is_enum = self._is_enum_class(pc)

        base_class_stereotypes = {"enum", "pydantic", "abstract"}
        for dec in pc.decorators:
            if dec.lower() in base_class_stereotypes:
                continue
            lines.append(f"@{dec}")

        bases = list(pc.bases)
        if pc.implements:
            bases.extend(pc.implements)
        if is_enum and "Enum" not in bases:
            bases.append("Enum")
        if (
            any("pydantic" in s.lower() for s in pc.stereotypes + pc.decorators)
            and "BaseModel" not in bases
        ):
            bases.append("BaseModel")
        if (
            any("abstract" in s.lower() for s in pc.stereotypes + pc.decorators)
            and "ABC" not in bases
        ):
            bases.append("ABC")
        bases_str = f"({', '.join(bases)})" if bases else ""

        lines.append(f"class {pc.name}{bases_str}:")

        has_body = False
        display_stereotypes = [s for s in pc.stereotypes if s.lower() != "enum"]
        doc_parts = []
        if display_stereotypes:
            doc_parts.append(f"Stereotypes: {', '.join(display_stereotypes)}")
        if pc.compositions:
            doc_parts.append(f"Composes: {', '.join(t for t, _ in pc.compositions)}")
        if pc.aggregations:
            doc_parts.append(f"Aggregates: {', '.join(t for t, _ in pc.aggregations)}")
        if pc.dependencies:
            doc_parts.append(f"Depends on: {', '.join(t for t, _ in pc.dependencies)}")
        if doc_parts:
            lines.append(f'    """{"; ".join(doc_parts)}"""')
            has_body = True

        if is_enum:
            for var in pc.variables:
                lines.append(f'    {var.name} = "{var.name}"')
                has_body = True
        else:
            for var in pc.variables:
                ann = f": {var.type_annotation}" if var.type_annotation else ""
                default = f" = {var.default}" if var.default else ""
                lines.append(f"    {var.name}{ann}{default}")
                has_body = True

        for method in pc.methods:
            lines.extend(self._generate_method(method))
            has_body = True

        if not has_body:
            lines.append("    pass")

        return lines

    def _generate_method(self, method: FunctionInfo) -> list[str]:
        extra = getattr(method, "_extra", {})
        lines = [f"    @{dec}" for dec in method.decorators]

        async_kw = "async " if method.is_async else ""

        params = []
        is_method = extra.get("is_method", False)
        is_classmethod = extra.get("is_classmethod", False)
        is_staticmethod = extra.get("is_staticmethod", False)
        if is_classmethod:
            params.append("cls")
        elif is_method and not is_staticmethod:
            params.append("self")

        for p in method.parameters:
            p_str = p.name
            if p.type_annotation:
                p_str += f": {p.type_annotation}"
            if p.default:
                p_str += f" = {p.default}"
            params.append(p_str)

        params_str = ", ".join(params)
        returns = f" -> {method.return_type}" if method.return_type else ""

        lines.append(f"    {async_kw}def {method.name}({params_str}){returns}:")
        lines.append("        ...")
        return lines

    def _generate_file_code(
        self,
        classes: list[ParsedClass],
        namespace: str,
        class_to_file: dict[str, str],
    ) -> str:
        lines: list[str] = ["from __future__ import annotations", ""]
        import_lines = self._generate_imports(classes)
        if import_lines:
            lines.extend(import_lines)
            lines.append("")

        local_class_names = {pc.name for pc in classes}
        needed_imports: dict[str, set[str]] = {}
        for pc in classes:
            refs = self._extract_type_references(pc)
            for ref_name in refs:
                if ref_name in class_to_file and ref_name not in local_class_names:
                    needed_imports.setdefault(class_to_file[ref_name], set()).add(
                        ref_name
                    )

        for target_file in sorted(needed_imports):
            names = ", ".join(sorted(needed_imports[target_file]))
            lines.append(f"from {namespace}.{target_file} import {names}")
        if needed_imports:
            lines.append("")

        for pc in classes:
            lines.extend(self._generate_class(pc))
            lines.append("")

        return "\n".join(lines)

    def _generate_init_code(self, namespace: str, class_to_file: dict[str, str]) -> str:
        lines: list[str] = [
            f'"""{namespace.split(".")[-1]} package."""',
            "",
        ]
        file_to_classes: dict[str, list[str]] = {}
        for cls_name, file_stem in sorted(class_to_file.items()):
            file_to_classes.setdefault(file_stem, []).append(cls_name)
        for file_stem in sorted(file_to_classes):
            names = ", ".join(sorted(file_to_classes[file_stem]))
            lines.append(f"from {namespace}.{file_stem} import {names}")
        lines.append("")
        all_names = sorted(class_to_file.keys())
        all_str = ", ".join(f'"{n}"' for n in all_names)
        lines.append(f"__all__ = [{all_str}]")
        lines.append("")
        return "\n".join(lines)

    def write_modules(
        self, parsed: list[ParsedClass], output_dir: str
    ) -> dict[str, str]:
        groups: OrderedDict[str, list[ParsedClass]] = OrderedDict()
        for pc in parsed:
            ns = pc.namespace or ""
            groups.setdefault(ns, []).append(pc)

        written: dict[str, str] = {}
        all_intermediate_dirs: set[str] = set()

        for ns, group in groups.items():
            dir_path_rel = ns.replace(".", "/") if ns else "scaffold"
            full_dir = os.path.join(output_dir, dir_path_rel)
            os.makedirs(full_dir, exist_ok=True)

            parts = dir_path_rel.split("/")
            for i in range(1, len(parts)):
                parent = os.path.join(output_dir, *parts[:i])
                all_intermediate_dirs.add(parent)

            enums = [pc for pc in group if self._is_enum_class(pc)]
            non_enums = [pc for pc in group if not self._is_enum_class(pc)]

            class_to_file: dict[str, str] = {}
            for pc in enums:
                class_to_file[pc.name] = "types"
            for pc in non_enums:
                class_to_file[pc.name] = self._class_name_to_filename(pc.name)

            if enums:
                code = self._generate_file_code(enums, ns, class_to_file)
                fpath = os.path.join(full_dir, "types.py")
                with open(fpath, "w", encoding="utf-8") as f:
                    f.write(code)
                written[fpath] = code

            for pc in non_enums:
                file_stem = class_to_file[pc.name]
                code = self._generate_file_code([pc], ns, class_to_file)
                fpath = os.path.join(full_dir, f"{file_stem}.py")
                with open(fpath, "w", encoding="utf-8") as f:
                    f.write(code)
                written[fpath] = code

            init_code = self._generate_init_code(ns, class_to_file)
            init_path = os.path.join(full_dir, "__init__.py")
            with open(init_path, "w", encoding="utf-8") as f:
                f.write(init_code)
            written[init_path] = init_code

        for dir_path in sorted(all_intermediate_dirs):
            os.makedirs(dir_path, exist_ok=True)
            init_path = os.path.join(dir_path, "__init__.py")
            if init_path not in written:
                pkg_name = os.path.basename(dir_path)
                with open(init_path, "w", encoding="utf-8") as f:
                    f.write(f'"""{pkg_name} package."""\n')
                written[init_path] = f'"""{pkg_name} package."""\n'

        return written