Update entity_relationship_generator.py

entity_relationship_generator.py

@@ -213,18 +213,30 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                 'splines': 'ortho',
                 'bgcolor': 'white',
                 'pad': '0.5',
-                'nodesep': '2.0',
-                'ranksep': '2.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'
             }
         )
         
-        base_color = '#19191a'
-        lightening_factor = 0.15
-        
         entities = data.get('entities', [])
         relationships = data.get('relationships', [])
         
-        for i, entity in enumerate(entities):
+        for entity in entities:
             entity_name = entity.get('name')
             entity_type = entity.get('type', 'strong')
             attributes = entity.get('attributes', [])
@@ -232,49 +244,28 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             if not entity_name:
                 raise ValueError(f"Invalid entity: {entity}")
             
-            current_depth = i % 6
-            
-            if not isinstance(base_color, str) or not base_color.startswith('#') or len(base_color) != 7:
-                base_color_safe = '#19191a'
-            else:
-                base_color_safe = base_color
-                
-            base_r = int(base_color_safe[1:3], 16)
-            base_g = int(base_color_safe[3:5], 16)
-            base_b = int(base_color_safe[5:7], 16)
-            
-            current_r = base_r + int((255 - base_r) * current_depth * lightening_factor)
-            current_g = base_g + int((255 - base_g) * current_depth * lightening_factor)
-            current_b = base_b + int((255 - base_b) * current_depth * lightening_factor)
-            
-            current_r = min(255, current_r)
-            current_g = min(255, current_g)
-            current_b = min(255, current_b)
-            
-            node_color = f'#{current_r:02x}{current_g:02x}{current_b:02x}'
-            font_color = 'white' if current_depth * lightening_factor < 0.6 else 'black'
-            
-            entity_label = f"{entity_name}\\n"
+            entity_label = f"<B>{entity_name}</B>"
             
             if attributes:
+                entity_label += "|"
                 primary_keys = []
                 foreign_keys = []
                 regular_attrs = []
                 
                 for attr in attributes:
                     attr_name = attr.get('name', '')
-                    attr_type = attr.get('type', 'key')
+                    attr_type = attr.get('type', 'regular')
                     is_multivalued = attr.get('multivalued', False)
                     is_derived = attr.get('derived', False)
                     is_composite = attr.get('composite', False)
                     
                     if attr_type == 'primary_key':
                         if is_multivalued:
-                            primary_keys.append(f"{{{{ {attr_name} }}}}")
+                            primary_keys.append(f"<U>{{{{ {attr_name} }}}}</U>")
                         else:
-                            primary_keys.append(f"[PK] {attr_name}")
+                            primary_keys.append(f"<U>{attr_name}</U>")
                     elif attr_type == 'foreign_key':
-                        foreign_keys.append(f"[FK] {attr_name}")
+                        foreign_keys.append(f"<I>{attr_name}</I> (FK)")
                     else:
                         attr_display = attr_name
                         if is_derived:
@@ -285,29 +276,27 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                             attr_display = f"( {attr_display} )"
                         regular_attrs.append(attr_display)
                 
-                if primary_keys:
-                    entity_label += "\\n".join(primary_keys) + "\\n"
-                if foreign_keys:
-                    entity_label += "\\n".join(foreign_keys) + "\\n"
-                if regular_attrs:
-                    entity_label += "\\n".join(regular_attrs)
+                all_attrs = primary_keys + foreign_keys + regular_attrs
+                entity_label += "\\l".join(all_attrs) + "\\l"
             
             if entity_type == 'weak':
-                shape = 'doubleoctagon'
+                shape = 'record'
                 style = 'filled'
+                fillcolor = '#f0f0f0'
+                penwidth = '3'
             else:
-                shape = 'box'
-                style = 'filled,rounded'
+                shape = 'record'
+                style = 'filled'
+                fillcolor = 'white'
+                penwidth = '1.5'
             
             dot.node(
                 entity_name,
-                entity_label,
+                f"<{entity_label}>",
                 shape=shape,
                 style=style,
-                fillcolor=node_color,
-                fontcolor=font_color,
-                fontsize='10',
-                fontname='Helvetica'
+                fillcolor=fillcolor,
+                penwidth=penwidth
             )
         
         for relationship in relationships:
@@ -319,48 +308,41 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             if not rel_name or len(entities_involved) < 2:
                 raise ValueError(f"Invalid relationship: {relationship}")
             
-            rel_node_color = '#e6f3ff'
-            
             if rel_type == 'identifying':
                 rel_shape = 'diamond'
-                rel_style = 'filled,bold'
-                rel_color = '#4a90e2'
+                rel_style = 'filled'
+                rel_color = '#d0d0d0'
+                rel_penwidth = '3'
             elif rel_type == 'weak':
                 rel_shape = 'diamond'
-                rel_style = 'filled,dashed'
-                rel_color = '#a0a0a0'
+                rel_style = 'filled'
+                rel_color = '#e8e8e8'
+                rel_penwidth = '2'
             else:
                 rel_shape = 'diamond'
                 rel_style = 'filled'
-                rel_color = '#4a90e2'
+                rel_color = '#d0d0d0'
+                rel_penwidth = '1.5'
             
             dot.node(
                 rel_name,
-                rel_name,
+                f"<B>{rel_name}</B>",
                 shape=rel_shape,
                 style=rel_style,
                 fillcolor=rel_color,
-                fontcolor='white',
-                fontsize='10',
-                fontname='Helvetica'
+                fontcolor='black',
+                penwidth=rel_penwidth
             )
             
             for entity in entities_involved:
                 cardinality = cardinalities.get(entity, '1')
                 
-                edge_label = cardinality
-                if cardinality in ['1:1', '1:N', 'M:N', '1', 'N', 'M']:
-                    pass
-                else:
-                    edge_label = cardinality
-                
                 dot.edge(
                     entity,
                     rel_name,
-                    label=edge_label,
-                    color='#4a4a4a',
-                    fontsize='9',
-                    fontcolor='#4a4a4a'
+                    label=cardinality,
+                    labelfontsize='10',
+                    labelfontcolor='black'
                 )
 
         with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp: