Update entity_relationship_generator.py

entity_relationship_generator.py

@@ -214,22 +214,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                 'bgcolor': 'white',
                 'pad': '0.5',
                 'nodesep': '1.2',
-                'ranksep': '1.8',
-                'fontname': 'Arial',
-                'dpi': '300',
-                'overlap': 'false'
-            },
-            node_attr={
-                'fontname': 'Arial',
-                'fontsize': '10',
-                'color': 'black',
-                'penwidth': '1.5'
-            },
-            edge_attr={
-                'fontname': 'Arial',
-                'fontsize': '9',
-                'color': 'black',
-                'penwidth': '1'
+                'ranksep': '1.8'
             }
         )
         
@@ -244,14 +229,10 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             if not entity_name:
                 raise ValueError(f"Invalid entity: {entity}")
             
-            entity_label = f"{entity_name}"
+            entity_parts = [entity_name]
             
             if attributes:
-                entity_label += "|"
-                primary_keys = []
-                foreign_keys = []
-                regular_attrs = []
-                
+                attr_section = ""
                 for attr in attributes:
                     attr_name = attr.get('name', '')
                     attr_type = attr.get('type', 'regular')
@@ -261,11 +242,11 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                     
                     if attr_type == 'primary_key':
                         if is_multivalued:
-                            primary_keys.append(f"PK: {{ {attr_name} }}")
+                            attr_display = f"PK: {{ {attr_name} }}"
                         else:
-                            primary_keys.append(f"PK: {attr_name}")
+                            attr_display = f"PK: {attr_name}"
                     elif attr_type == 'foreign_key':
-                        foreign_keys.append(f"FK: {attr_name}")
+                        attr_display = f"FK: {attr_name}"
                     else:
                         attr_display = attr_name
                         if is_derived:
@@ -274,21 +255,27 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                             attr_display = f"{{ {attr_display} }}"
                         if is_composite:
                             attr_display = f"( {attr_display} )"
-                        regular_attrs.append(attr_display)
+                    
+                    if attr_section:
+                        attr_section += "\\l"
+                    attr_section += attr_display
                 
-                all_attrs = primary_keys + foreign_keys + regular_attrs
-                entity_label += "\\l".join(all_attrs) + "\\l"
+                if attr_section:
+                    attr_section += "\\l"
+                    entity_parts.append(attr_section)
+            
+            entity_label = "|".join(entity_parts)
             
             if entity_type == 'weak':
                 shape = 'record'
                 style = 'filled'
-                fillcolor = '#f0f0f0'
-                penwidth = '3'
+                fillcolor = '#e8e8e8'
+                penwidth = '2'
             else:
                 shape = 'record'
                 style = 'filled'
-                fillcolor = 'white'
-                penwidth = '1.5'
+                fillcolor = '#d0d0d0'
+                penwidth = '1'
             
             dot.node(
                 entity_name,
@@ -296,6 +283,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                 shape=shape,
                 style=style,
                 fillcolor=fillcolor,
+                color='black',
                 penwidth=penwidth
             )
         
@@ -311,18 +299,18 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             if rel_type == 'identifying':
                 rel_shape = 'diamond'
                 rel_style = 'filled'
-                rel_color = '#d0d0d0'
-                rel_penwidth = '3'
+                rel_color = '#c0c0c0'
+                rel_penwidth = '2'
             elif rel_type == 'weak':
                 rel_shape = 'diamond'
                 rel_style = 'filled'
-                rel_color = '#e8e8e8'
-                rel_penwidth = '2'
+                rel_color = '#e0e0e0'
+                rel_penwidth = '1'
             else:
                 rel_shape = 'diamond'
                 rel_style = 'filled'
-                rel_color = '#d0d0d0'
-                rel_penwidth = '1.5'
+                rel_color = '#c0c0c0'
+                rel_penwidth = '1'
             
             dot.node(
                 rel_name,
@@ -330,6 +318,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                 shape=rel_shape,
                 style=rel_style,
                 fillcolor=rel_color,
+                color='black',
                 fontcolor='black',
                 penwidth=rel_penwidth
             )
@@ -341,8 +330,7 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                     entity,
                     rel_name,
                     label=cardinality,
-                    labelfontsize='10',
-                    labelfontcolor='black'
+                    color='black'
                 )
 
         with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp: