Diagram Creation

class EntityRelationshipDiagram:
    """Holds information about an entity relationship diagram for a set of data model classes and
    their relationships, and provides methods to render the diagram.
    
    Attributes:
        models (SortedDict[str, ModelInfo]): Mapping of ModelInfo instances for models included
            in the diagram. Each key is the string representation of the fully qualified name of
            the model.
        edges (SortedDict[str, Edge]): Mapping of edges representing relationships between the models.
    """
    
    def __init__(self) -> None:
        """Initialize empty EntityRelationshipDiagram."""
        
    def add_model(self, model: type, recurse: bool = True) -> None:
        """Add a data model class to the diagram.

        Args:
            model (type): Data model class to add to the diagram.
            recurse (bool, optional): Whether to recursively add models referenced by fields of
                the given model. Defaults to True.

        Raises:
            UnknownModelTypeError: If the model is not recognized as a data model class type that
                is supported by registered plugins.
            UnevaluatedForwardRefError: If the model contains a forward reference that cannot be
                automatically resolved.
        """
        
    def draw(self, out: Union[str, os.PathLike], graph_attr: Optional[Mapping[str, Any]] = None,
             node_attr: Optional[Mapping[str, Any]] = None, edge_attr: Optional[Mapping[str, Any]] = None,
             **kwargs) -> None:
        """Render entity relationship diagram for given data model classes to file. The file format
        can be inferred from the file extension. Typical formats include '.png', '.svg', and
        '.pdf'.

        Args:
            out (str | os.PathLike): Output file path for rendered diagram.
            graph_attr (Mapping[str, Any] | None, optional): Override any graph attributes on
                the `pygraphviz.AGraph` instance. Defaults to None.
            node_attr (Mapping[str, Any] | None, optional): Override any node attributes for all
                nodes on the `pygraphviz.AGraph` instance. Defaults to None.
            edge_attr (Mapping[str, Any] | None, optional): Override any edge attributes for all
                edges on the `pygraphviz.AGraph` instance. Defaults to None.
            **kwargs: Additional keyword arguments to
                [`pygraphviz.AGraph.draw`][pygraphviz.AGraph.draw].
        """
        
    def to_graphviz(self, graph_attr: Optional[Mapping[str, Any]] = None,
                    node_attr: Optional[Mapping[str, Any]] = None, 
                    edge_attr: Optional[Mapping[str, Any]] = None) -> pgv.AGraph:
        """Return [`pygraphviz.AGraph`][pygraphviz.agraph.AGraph] instance for diagram.

        Args:
            graph_attr (Mapping[str, Any] | None, optional): Override any graph attributes on
                the `pygraphviz.AGraph` instance. Defaults to None.
            node_attr (Mapping[str, Any] | None, optional): Override any node attributes for all
                nodes on the `pygraphviz.AGraph` instance. Defaults to None.
            edge_attr (Mapping[str, Any] | None, optional): Override any edge attributes for all
                edges on the `pygraphviz.AGraph` instance. Defaults to None.

        Returns:
            pygraphviz.AGraph: graph object for diagram
        """
        
    def to_dot(self, graph_attr: Optional[Mapping[str, Any]] = None,
               node_attr: Optional[Mapping[str, Any]] = None,
               edge_attr: Optional[Mapping[str, Any]] = None) -> str:
        """Generate Graphviz [DOT language](https://graphviz.org/doc/info/lang.html) representation
        of entity relationship diagram for given data model classes.

        Args:
            graph_attr (Mapping[str, Any] | None, optional): Override any graph attributes on
                the `pygraphviz.AGraph` instance. Defaults to None.
            node_attr (Mapping[str, Any] | None, optional): Override any node attributes for all
                nodes on the `pygraphviz.AGraph` instance. Defaults to None.
            edge_attr (Mapping[str, Any] | None, optional): Override any edge attributes for all
                edges on the `pygraphviz.AGraph` instance. Defaults to None.

        Returns:
            str: DOT language representation of diagram
        """

from erdantic import EntityRelationshipDiagram
from erdantic.core import EntityRelationshipDiagram  # Alternative import
import os
from typing import Union, Optional, Mapping, Any

from pydantic import BaseModel
from erdantic import EntityRelationshipDiagram

class User(BaseModel):
    name: str
    email: str

class Post(BaseModel):
    title: str
    author: User

# Create diagram and add models
diagram = EntityRelationshipDiagram()
diagram.add_model(User)
diagram.add_model(Post)

# Render to PNG file
diagram.draw("models.png")

# Custom styling with Graphviz attributes
graph_attrs = {
    "rankdir": "TD",  # Top-down layout
    "bgcolor": "white",
    "label": "My Custom ERD"
}

node_attrs = {
    "fontname": "Arial",
    "fontsize": 12,
    "fillcolor": "lightblue",
    "style": "filled"
}

edge_attrs = {
    "color": "blue",
    "penwidth": 2
}

diagram.draw(
    "custom_styled.svg", 
    graph_attr=graph_attrs,
    node_attr=node_attrs, 
    edge_attr=edge_attrs
)

# Add models without recursing into their field types
diagram.add_model(User, recurse=False)
diagram.add_model(Post, recurse=False)

# This creates a diagram with just User and Post,
# without analyzing their field relationships

# Generate DOT language representation
dot_string = diagram.to_dot()
print(dot_string)

# Save DOT to file for external processing
with open("diagram.dot", "w") as f:
    f.write(dot_string)

# Access collected model information
for model_key, model_info in diagram.models.items():
    print(f"Model: {model_info.name}")
    for field_name, field_info in model_info.fields.items():
        print(f"  Field: {field_name} ({field_info.type_name})")

# Access relationship information
for edge_key, edge in diagram.edges.items():
    print(f"Relationship: {edge.source_model_full_name}.{edge.source_field_name} -> {edge.target_model_full_name}")
    print(f"  Cardinality: {edge.target_cardinality.value}")
    print(f"  Modality: {edge.target_modality.value}")

# Default graph attributes
DEFAULT_GRAPH_ATTR = (
    ("nodesep", "0.5"),
    ("ranksep", "1.5"), 
    ("rankdir", "LR"),  # Left-to-right layout
    ("label", f"Created by erdantic v{__version__} <https://github.com/drivendataorg/erdantic>"),
    ("fontname", "Times New Roman,Times,Liberation Serif,serif"),
    ("fontsize", "9"),
    ("fontcolor", "gray66"),
)

# Default node attributes
DEFAULT_NODE_ATTR = (
    ("fontname", "Times New Roman,Times,Liberation Serif,serif"),
    ("fontsize", 14),
    ("shape", "plain"),  # Uses HTML table format
)

# Default edge attributes  
DEFAULT_EDGE_ATTR = (("dir", "both"),)  # Bidirectional arrows

from erdantic.core import DEFAULT_GRAPH_ATTR, DEFAULT_NODE_ATTR, DEFAULT_EDGE_ATTR

# Use defaults as base for custom styling
custom_graph_attr = dict(DEFAULT_GRAPH_ATTR)
custom_graph_attr.update({"bgcolor": "lightgray", "rankdir": "TB"})

diagram.draw("custom.png", graph_attr=custom_graph_attr)

from erdantic.exceptions import UnknownModelTypeError, UnevaluatedForwardRefError

try:
    diagram.add_model(MyModel)
except UnknownModelTypeError as e:
    print(f"Model type not supported: {e.model}")
    print(f"Available plugins: {e.available_plugins}")
except UnevaluatedForwardRefError as e:
    print(f"Cannot resolve forward reference '{e.name}' in model {e.model_full_name}")