Graphify / entity_relationship_generator.py
ZahirJS's picture
Update entity_relationship_generator.py
0ff231b verified
raw
history blame
8.62 kB
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)}"