entity_relationship_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os

def generate_entity_relationship_diagram(json_input: str, output_format: str) -> str:
    try:
        if not json_input.strip():
            return "Error: Empty input"
            
        data = json.loads(json_input)
        
        if 'entities' not in data:
            raise ValueError("Missing required field: entities")

        dot = graphviz.Graph(comment='ER Diagram', engine='neato')
        dot.attr(
            bgcolor='white',
            pad='1.5',
            overlap='false',
            splines='true',
            sep='+25',
            esep='+15'
        )
        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', [])
        
        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':
                dot.node(
                    entity_name, 
                    entity_name, 
                    shape='box', 
                    style='filled', 
                    fillcolor='#f0f0f0',
                    color='#404040',
                    penwidth='3',
                    width='1.8',
                    height='0.8'
                )
            else:
                dot.node(
                    entity_name, 
                    entity_name, 
                    shape='box', 
                    style='filled', 
                    fillcolor='#e8e8e8',
                    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}"
                
                # CAMBIO CLAVE: Removemos HTML markup que puede causar problemas
                if attr_type == 'primary_key':
                    dot.node(
                        attr_id, 
                        f'{attr_name} (PK)', 
                        shape='ellipse', 
                        style='filled',
                        fillcolor='#d8d8d8',
                        color='#404040',
                        width='1.2',
                        height='0.6'
                    )
                elif attr_type == 'partial_key':
                    dot.node(
                        attr_id, 
                        f'{attr_name} (Partial)', 
                        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')
        
        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:
                continue
            
            if rel_type == 'isa':
                parent = relationship.get('parent')
                children = relationship.get('children', [])
                
                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'
                    )

        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}"

    except json.JSONDecodeError:
        return "Error: Invalid JSON format"
    except Exception as e:
        return f"Error: {str(e)}"