0
# Diagram Creation
1

2
The `EntityRelationshipDiagram` class is the core component for creating and managing entity relationship diagrams programmatically. It provides fine-grained control over diagram construction, model analysis, and rendering options.
3

4
## EntityRelationshipDiagram Class
5

6
```python { .api }
7
class EntityRelationshipDiagram:
8
    """Holds information about an entity relationship diagram for a set of data model classes and
9
    their relationships, and provides methods to render the diagram.
10
    
11
    Attributes:
12
        models (SortedDict[str, ModelInfo]): Mapping of ModelInfo instances for models included
13
            in the diagram. Each key is the string representation of the fully qualified name of
14
            the model.
15
        edges (SortedDict[str, Edge]): Mapping of edges representing relationships between the models.
16
    """
17
    
18
    def __init__(self) -> None:
19
        """Initialize empty EntityRelationshipDiagram."""
20
        
21
    def add_model(self, model: type, recurse: bool = True) -> None:
22
        """Add a data model class to the diagram.
23

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

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

43
        Args:
44
            out (str | os.PathLike): Output file path for rendered diagram.
45
            graph_attr (Mapping[str, Any] | None, optional): Override any graph attributes on
46
                the `pygraphviz.AGraph` instance. Defaults to None.
47
            node_attr (Mapping[str, Any] | None, optional): Override any node attributes for all
48
                nodes on the `pygraphviz.AGraph` instance. Defaults to None.
49
            edge_attr (Mapping[str, Any] | None, optional): Override any edge attributes for all
50
                edges on the `pygraphviz.AGraph` instance. Defaults to None.
51
            **kwargs: Additional keyword arguments to
52
                [`pygraphviz.AGraph.draw`][pygraphviz.AGraph.draw].
53
        """
54
        
55
    def to_graphviz(self, graph_attr: Optional[Mapping[str, Any]] = None,
56
                    node_attr: Optional[Mapping[str, Any]] = None, 
57
                    edge_attr: Optional[Mapping[str, Any]] = None) -> pgv.AGraph:
58
        """Return [`pygraphviz.AGraph`][pygraphviz.agraph.AGraph] instance for diagram.
59

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

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

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

86
        Returns:
87
            str: DOT language representation of diagram
88
        """
89
```
90

91
## Required Imports
92

93
```python
94
from erdantic import EntityRelationshipDiagram
95
from erdantic.core import EntityRelationshipDiagram  # Alternative import
96
import os
97
from typing import Union, Optional, Mapping, Any
98
```
99

100
## Usage Examples
101

102
### Basic Diagram Creation
103

104
```python
105
from pydantic import BaseModel
106
from erdantic import EntityRelationshipDiagram
107

108
class User(BaseModel):
109
    name: str
110
    email: str
111

112
class Post(BaseModel):
113
    title: str
114
    author: User
115

116
# Create diagram and add models
117
diagram = EntityRelationshipDiagram()
118
diagram.add_model(User)
119
diagram.add_model(Post)
120

121
# Render to PNG file
122
diagram.draw("models.png")
123
```
124

125
### Advanced Configuration
126

127
```python
128
# Custom styling with Graphviz attributes
129
graph_attrs = {
130
    "rankdir": "TD",  # Top-down layout
131
    "bgcolor": "white",
132
    "label": "My Custom ERD"
133
}
134

135
node_attrs = {
136
    "fontname": "Arial",
137
    "fontsize": 12,
138
    "fillcolor": "lightblue",
139
    "style": "filled"
140
}
141

142
edge_attrs = {
143
    "color": "blue",
144
    "penwidth": 2
145
}
146

147
diagram.draw(
148
    "custom_styled.svg", 
149
    graph_attr=graph_attrs,
150
    node_attr=node_attrs, 
151
    edge_attr=edge_attrs
152
)
153
```
154

155
### Controlling Recursion
156

157
```python
158
# Add models without recursing into their field types
159
diagram.add_model(User, recurse=False)
160
diagram.add_model(Post, recurse=False)
161

162
# This creates a diagram with just User and Post,
163
# without analyzing their field relationships
164
```
165

166
### Working with DOT Format
167

168
```python
169
# Generate DOT language representation
170
dot_string = diagram.to_dot()
171
print(dot_string)
172

173
# Save DOT to file for external processing
174
with open("diagram.dot", "w") as f:
175
    f.write(dot_string)
176
```
177

178
### Accessing Diagram Data
179

180
```python
181
# Access collected model information
182
for model_key, model_info in diagram.models.items():
183
    print(f"Model: {model_info.name}")
184
    for field_name, field_info in model_info.fields.items():
185
        print(f"  Field: {field_name} ({field_info.type_name})")
186

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

194
## Default Graphviz Attributes
195

196
Erdantic provides sensible defaults for diagram appearance:
197

198
```python { .api }
199
# Default graph attributes
200
DEFAULT_GRAPH_ATTR = (
201
    ("nodesep", "0.5"),
202
    ("ranksep", "1.5"), 
203
    ("rankdir", "LR"),  # Left-to-right layout
204
    ("label", f"Created by erdantic v{__version__} <https://github.com/drivendataorg/erdantic>"),
205
    ("fontname", "Times New Roman,Times,Liberation Serif,serif"),
206
    ("fontsize", "9"),
207
    ("fontcolor", "gray66"),
208
)
209

210
# Default node attributes
211
DEFAULT_NODE_ATTR = (
212
    ("fontname", "Times New Roman,Times,Liberation Serif,serif"),
213
    ("fontsize", 14),
214
    ("shape", "plain"),  # Uses HTML table format
215
)
216

217
# Default edge attributes  
218
DEFAULT_EDGE_ATTR = (("dir", "both"),)  # Bidirectional arrows
219
```
220

221
These constants can be imported and used as reference or starting points for custom styling:
222

223
```python
224
from erdantic.core import DEFAULT_GRAPH_ATTR, DEFAULT_NODE_ATTR, DEFAULT_EDGE_ATTR
225

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

230
diagram.draw("custom.png", graph_attr=custom_graph_attr)
231
```
232

233
## Error Handling
234

235
```python
236
from erdantic.exceptions import UnknownModelTypeError, UnevaluatedForwardRefError
237

238
try:
239
    diagram.add_model(MyModel)
240
except UnknownModelTypeError as e:
241
    print(f"Model type not supported: {e.model}")
242
    print(f"Available plugins: {e.available_plugins}")
243
except UnevaluatedForwardRefError as e:
244
    print(f"Cannot resolve forward reference '{e.name}' in model {e.model_full_name}")
245
```