Update entity_relationship_generator.py

entity_relationship_generator.py

@@ -205,18 +205,10 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
         if 'entities' not in data:
             raise ValueError("Missing required field: entities")
 
-        dot = graphviz.Graph(
-            name='ERDiagram',
-            format='png',
-            graph_attr={
-                'rankdir': 'TB',
-                'splines': 'ortho',
-                'bgcolor': 'white',
-                'pad': '0.5',
-                'nodesep': '1.2',
-                'ranksep': '1.8'
-            }
-        )
+        dot = graphviz.Graph(comment='ER Diagram', engine='dot')
+        dot.attr(rankdir='TB', bgcolor='white', pad='0.5')
+        dot.attr('node', fontname='Arial', fontsize='10', color='black')
+        dot.attr('edge', fontname='Arial', fontsize='9', color='black')
         
         entities = data.get('entities', [])
         relationships = data.get('relationships', [])
@@ -227,12 +219,12 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             attributes = entity.get('attributes', [])
             
             if not entity_name:
-                raise ValueError(f"Invalid entity: {entity}")
+                continue
             
-            entity_parts = [entity_name]
+            label_parts = [entity_name]
             
             if attributes:
-                attr_section = ""
+                attr_lines = []
                 for attr in attributes:
                     attr_name = attr.get('name', '')
                     attr_type = attr.get('type', 'regular')
@@ -256,36 +248,23 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         if is_composite:
                             attr_display = f"( {attr_display} )"
                     
-                    if attr_section:
-                        attr_section += "\\l"
-                    attr_section += attr_display
+                    attr_lines.append(attr_display)
                 
-                if attr_section:
-                    attr_section += "\\l"
-                    entity_parts.append(attr_section)
+                if attr_lines:
+                    label_parts.extend(attr_lines)
             
-            entity_label = "|".join(entity_parts)
+            label = "\\n".join(label_parts)
             
             if entity_type == 'weak':
-                shape = 'record'
                 style = 'filled'
                 fillcolor = '#e8e8e8'
                 penwidth = '2'
             else:
-                shape = 'record'
                 style = 'filled'
                 fillcolor = '#d0d0d0'
                 penwidth = '1'
             
-            dot.node(
-                entity_name,
-                entity_label,
-                shape=shape,
-                style=style,
-                fillcolor=fillcolor,
-                color='black',
-                penwidth=penwidth
-            )
+            dot.node(entity_name, label, shape='box', style=style, fillcolor=fillcolor, penwidth=penwidth)
         
         for relationship in relationships:
             rel_name = relationship.get('name')
@@ -294,44 +273,22 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             cardinalities = relationship.get('cardinalities', {})
             
             if not rel_name or len(entities_involved) < 2:
-                raise ValueError(f"Invalid relationship: {relationship}")
+                continue
             
             if rel_type == 'identifying':
-                rel_shape = 'diamond'
-                rel_style = 'filled'
-                rel_color = '#c0c0c0'
-                rel_penwidth = '2'
-            elif rel_type == 'weak':
-                rel_shape = 'diamond'
-                rel_style = 'filled'
-                rel_color = '#e0e0e0'
-                rel_penwidth = '1'
+                style = 'filled'
+                fillcolor = '#c0c0c0'
+                penwidth = '2'
             else:
-                rel_shape = 'diamond'
-                rel_style = 'filled'
-                rel_color = '#c0c0c0'
-                rel_penwidth = '1'
+                style = 'filled'
+                fillcolor = '#c0c0c0'
+                penwidth = '1'
             
-            dot.node(
-                rel_name,
-                rel_name,
-                shape=rel_shape,
-                style=rel_style,
-                fillcolor=rel_color,
-                color='black',
-                fontcolor='black',
-                penwidth=rel_penwidth
-            )
+            dot.node(rel_name, rel_name, shape='diamond', style=style, fillcolor=fillcolor, penwidth=penwidth)
             
             for entity in entities_involved:
                 cardinality = cardinalities.get(entity, '1')
-                
-                dot.edge(
-                    entity,
-                    rel_name,
-                    label=cardinality,
-                    color='black'
-                )
+                dot.edge(entity, rel_name, label=cardinality)
 
         with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
             dot.render(tmp.name, format=output_format, cleanup=True)