Update entity_relationship_generator.py

entity_relationship_generator.py

@@ -4,43 +4,6 @@ from tempfile import NamedTemporaryFile
 import os
 
 def generate_entity_relationship_diagram(json_input: str, output_format: str) -> str:
-    """
-    Generates an entity relationship diagram from JSON input.
-
-    Args:
-        json_input (str): A JSON string describing the entity relationship diagram structure.
-
-    Expected JSON Format Example:
-    {
-      "entities": [
-        {
-          "name": "CUSTOMER",
-          "type": "strong"
-        },
-        {
-          "name": "ORDER", 
-          "type": "strong"
-        },
-        {
-          "name": "REVIEW",
-          "type": "weak"
-        }
-      ],
-      "relationships": [
-        {
-          "name": "Places",
-          "entities": ["CUSTOMER", "ORDER"],
-          "cardinalities": {
-            "CUSTOMER": "1",
-            "ORDER": "N"
-          }
-        }
-      ]
-    }
-
-    Returns:
-        str: The filepath to the generated PNG image file.
-    """
     try:
         if not json_input.strip():
             return "Error: Empty input"
@@ -50,104 +13,217 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
         if 'entities' not in data:
             raise ValueError("Missing required field: entities")
 
-        # Configuración ultra simple - dejar que Graphviz organice libremente
-        dot = graphviz.Graph(comment='ER Diagram', engine='fdp')
+        dot = graphviz.Graph(comment='ER Diagram', engine='neato')
         dot.attr(
             bgcolor='white',
-            pad='1.0',
+            pad='1.5',
             overlap='false',
-            splines='line',
-            K='2'
+            splines='true',
+            sep='+25',
+            esep='+15'
         )
-        dot.attr('node', fontname='Arial', fontsize='11')
-        dot.attr('edge', fontname='Arial', fontsize='10')
+        dot.attr('node', fontname='Arial', fontsize='10', color='#404040')
+        dot.attr('edge', fontname='Arial', fontsize='9', color='#404040')
         
         entities = data.get('entities', [])
         relationships = data.get('relationships', [])
         
-        # 1. Crear todas las entidades como nodos simples
         for entity in entities:
             entity_name = entity.get('name')
             entity_type = entity.get('type', 'strong')
+            attributes = entity.get('attributes', [])
             
             if not entity_name:
                 continue
             
             if entity_type == 'weak':
-                # Entidad débil
                 dot.node(
                     entity_name, 
                     entity_name, 
                     shape='box', 
                     style='filled', 
-                    fillcolor='#e8e8e8',
+                    fillcolor='#f0f0f0',
                     color='#404040',
-                    penwidth='3'
+                    penwidth='3',
+                    width='1.8',
+                    height='0.8'
                 )
             else:
-                # Entidad fuerte
                 dot.node(
                     entity_name, 
                     entity_name, 
                     shape='box', 
                     style='filled', 
                     fillcolor='#e8e8e8',
-                    color='#404040'
-                )
-        
-        # 2. Crear todas las relaciones como nodos rombos
-        for relationship in relationships:
-            rel_name = relationship.get('name')
-            rel_type = relationship.get('type', 'regular')
-            
-            if not rel_name:
-                continue
-            
-            if rel_type == 'identifying':
-                # Relación identificadora
-                dot.node(
-                    rel_name, 
-                    rel_name, 
-                    shape='diamond', 
-                    style='filled',
-                    fillcolor='#d8d8d8',
                     color='#404040',
-                    penwidth='3'
-                )
-            else:
-                # Relación regular
-                dot.node(
-                    rel_name, 
-                    rel_name, 
-                    shape='diamond', 
-                    style='filled',
-                    fillcolor='#d8d8d8',
-                    color='#404040'
+                    penwidth='1',
+                    width='1.8',
+                    height='0.8'
                 )
+            
+            for i, attr in enumerate(attributes):
+                attr_name = attr.get('name', '')
+                attr_type = attr.get('type', 'regular')
+                
+                attr_id = f"{entity_name}_attr_{i}"
+                
+                if attr_type == 'primary_key':
+                    dot.node(
+                        attr_id, 
+                        f'<U><B>{attr_name}</B></U>', 
+                        shape='ellipse', 
+                        style='filled',
+                        fillcolor='#d8d8d8',
+                        color='#404040',
+                        width='1.2',
+                        height='0.6'
+                    )
+                elif attr_type == 'partial_key':
+                    dot.node(
+                        attr_id, 
+                        f'<U>{attr_name}</U>', 
+                        shape='ellipse', 
+                        style='filled,dashed',
+                        fillcolor='#d8d8d8',
+                        color='#404040',
+                        width='1.2',
+                        height='0.6'
+                    )
+                elif attr_type == 'multivalued':
+                    dot.node(
+                        attr_id, 
+                        attr_name, 
+                        shape='ellipse', 
+                        style='filled',
+                        fillcolor='#d8d8d8',
+                        color='#404040',
+                        penwidth='3',
+                        width='1.2',
+                        height='0.6'
+                    )
+                elif attr_type == 'derived':
+                    dot.node(
+                        attr_id, 
+                        f'/{attr_name}/', 
+                        shape='ellipse', 
+                        style='filled,dashed',
+                        fillcolor='#d8d8d8',
+                        color='#404040',
+                        width='1.2',
+                        height='0.6'
+                    )
+                elif attr_type == 'composite':
+                    dot.node(
+                        attr_id, 
+                        attr_name, 
+                        shape='ellipse', 
+                        style='filled',
+                        fillcolor='#d8d8d8',
+                        color='#404040',
+                        width='1.2',
+                        height='0.6'
+                    )
+                else:
+                    dot.node(
+                        attr_id, 
+                        attr_name, 
+                        shape='ellipse', 
+                        style='filled',
+                        fillcolor='#d8d8d8',
+                        color='#404040',
+                        width='1.2',
+                        height='0.6'
+                    )
+                
+                dot.edge(entity_name, attr_id, color='#404040', len='1.5')
         
-        # 3. Conectar entidades con relaciones (simple y directo)
         for relationship in relationships:
             rel_name = relationship.get('name')
+            rel_type = relationship.get('type', 'regular')
             entities_involved = relationship.get('entities', [])
             cardinalities = relationship.get('cardinalities', {})
+            rel_attributes = relationship.get('attributes', [])
             
-            if not rel_name or len(entities_involved) < 2:
+            if not rel_name:
                 continue
             
-            # Conectar cada entidad con la relación
-            for entity in entities_involved:
-                cardinality = cardinalities.get(entity, '1')
+            if rel_type == 'isa':
+                parent = relationship.get('parent')
+                children = relationship.get('children', [])
                 
-                # Conexión simple: entidad -> relación
-                dot.edge(
-                    entity, 
-                    rel_name, 
-                    label=cardinality,
-                    color='#404040',
-                    fontcolor='#000000'
-                )
+                if parent and children:
+                    isa_id = f"isa_{rel_name}"
+                    dot.node(
+                        isa_id, 
+                        'ISA', 
+                        shape='triangle', 
+                        style='filled',
+                        fillcolor='#d0d0d0',
+                        color='#404040',
+                        penwidth='2',
+                        width='1.0',
+                        height='0.8'
+                    )
+                    
+                    dot.edge(parent, isa_id, color='#404040', len='2.0')
+                    
+                    for child in children:
+                        dot.edge(isa_id, child, color='#404040', len='2.0')
+                
+            elif len(entities_involved) >= 2:
+                if rel_type == 'identifying':
+                    dot.node(
+                        rel_name, 
+                        rel_name, 
+                        shape='diamond', 
+                        style='filled',
+                        fillcolor='#c8c8c8',
+                        color='#404040',
+                        penwidth='3',
+                        width='1.8',
+                        height='1.0'
+                    )
+                else:
+                    dot.node(
+                        rel_name, 
+                        rel_name, 
+                        shape='diamond', 
+                        style='filled',
+                        fillcolor='#c8c8c8',
+                        color='#404040',
+                        penwidth='1',
+                        width='1.8',
+                        height='1.0'
+                    )
+                
+                for j, attr in enumerate(rel_attributes):
+                    attr_name = attr.get('name', '')
+                    attr_id = f"{rel_name}_attr_{j}"
+                    
+                    dot.node(
+                        attr_id, 
+                        attr_name, 
+                        shape='ellipse', 
+                        style='filled',
+                        fillcolor='#d8d8d8',
+                        color='#404040',
+                        width='1.0',
+                        height='0.5'
+                    )
+                    dot.edge(rel_name, attr_id, color='#404040', len='1.0')
+                
+                for entity in entities_involved:
+                    cardinality = cardinalities.get(entity, '1')
+                    dot.edge(
+                        entity, 
+                        rel_name, 
+                        label=f' {cardinality} ',
+                        color='#404040',
+                        len='2.5',
+                        fontcolor='#000000',
+                        fontsize='11'
+                    )
 
-        # Renderizar
         with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
             dot.render(tmp.name, format=output_format, cleanup=True)
             return f"{tmp.name}.{output_format}"