markitect-main/markitect/prompts/visualization/graph.py

"""
Graph visualization export for dependency graphs.

Exports DependencyGraph to DOT (Graphviz) and Mermaid diagram formats.
"""

from markitect.prompts.dependencies.models import DependencyGraph, EdgeType


# Edge type to style mappings
_DOT_EDGE_STYLES = {
    EdgeType.REQUIRES: 'style="solid"',
    EdgeType.GENERATES: 'style="dashed"',
    EdgeType.INCLUDES: 'style="dotted"',
}

_MERMAID_EDGE_STYLES = {
    EdgeType.REQUIRES: "-->",
    EdgeType.GENERATES: "-.->",
    EdgeType.INCLUDES: "==>",
}


class GraphExporter:
    """Export DependencyGraph to DOT and Mermaid formats."""

    @staticmethod
    def to_dot(graph: DependencyGraph, title: str = "Dependencies") -> str:
        """
        Export dependency graph to Graphviz DOT format.

        Args:
            graph: DependencyGraph to export
            title: Graph title

        Returns:
            DOT format string
        """
        lines = [
            f'digraph "{title}" {{',
            "  rankdir=LR;",
            f'  label="{title}";',
            "  node [shape=box];",
        ]

        # Add nodes
        for node in sorted(graph.nodes):
            safe_id = node.replace("-", "_")
            lines.append(f'  {safe_id} [label="{node}"];')

        # Add edges
        for source in sorted(graph.nodes):
            for target in sorted(graph.get_successors(source)):
                edge_type = graph.get_edge_type(source, target)
                style = _DOT_EDGE_STYLES.get(edge_type, 'style="solid"')
                safe_source = source.replace("-", "_")
                safe_target = target.replace("-", "_")
                label = edge_type.value if edge_type else "requires"
                lines.append(
                    f'  {safe_source} -> {safe_target} [{style} label="{label}"];'
                )

        lines.append("}")
        return "\n".join(lines)

    @staticmethod
    def to_mermaid(graph: DependencyGraph, title: str = "Dependencies") -> str:
        """
        Export dependency graph to Mermaid diagram format.

        Args:
            graph: DependencyGraph to export
            title: Graph title (used as comment)

        Returns:
            Mermaid format string
        """
        lines = [
            f"%%{{ title: {title} }}%%",
            "graph LR",
        ]

        # Add edges
        edges_added = False
        for source in sorted(graph.nodes):
            for target in sorted(graph.get_successors(source)):
                edge_type = graph.get_edge_type(source, target)
                arrow = _MERMAID_EDGE_STYLES.get(edge_type, "-->")
                label = edge_type.value if edge_type else "requires"
                lines.append(f"  {source}{arrow}|{label}|{target}")
                edges_added = True

        # Add isolated nodes (no edges)
        if not edges_added:
            for node in sorted(graph.nodes):
                lines.append(f"  {node}")
        else:
            # Add any isolated nodes that have no edges
            connected = set()
            for source in graph.nodes:
                if graph.get_successors(source) or graph.get_predecessors(source):
                    connected.add(source)
            for node in sorted(graph.nodes - connected):
                lines.append(f"  {node}")

        return "\n".join(lines)