Graphify / entity_relationship_generator.py
ZahirJS's picture
Update entity_relationship_generator.py
c0d2211 verified
raw
history blame
17.6 kB
# import graphviz
# import json
# 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.
# It must follow the Expected JSON Format Example below.
# Expected JSON Format Example:
# {
# "entities": [
# {
# "name": "Customer",
# "type": "strong",
# "attributes": [
# {"name": "customer_id", "type": "primary_key"},
# {"name": "first_name", "type": "regular"},
# {"name": "last_name", "type": "regular"},
# {"name": "email", "type": "regular"},
# {"name": "phone", "type": "multivalued", "multivalued": true},
# {"name": "age", "type": "derived", "derived": true},
# {"name": "address", "type": "composite", "composite": true}
# ]
# },
# {
# "name": "Order",
# "type": "strong",
# "attributes": [
# {"name": "order_id", "type": "primary_key"},
# {"name": "customer_id", "type": "foreign_key"},
# {"name": "order_date", "type": "regular"},
# {"name": "total_amount", "type": "regular"},
# {"name": "status", "type": "regular"},
# {"name": "shipping_address", "type": "composite", "composite": true}
# ]
# },
# {
# "name": "Product",
# "type": "strong",
# "attributes": [
# {"name": "product_id", "type": "primary_key"},
# {"name": "name", "type": "regular"},
# {"name": "description", "type": "regular"},
# {"name": "price", "type": "regular"},
# {"name": "category_id", "type": "foreign_key"},
# {"name": "stock_quantity", "type": "regular"}
# ]
# },
# {
# "name": "Category",
# "type": "strong",
# "attributes": [
# {"name": "category_id", "type": "primary_key"},
# {"name": "name", "type": "regular"},
# {"name": "description", "type": "regular"},
# {"name": "parent_category_id", "type": "foreign_key"}
# ]
# },
# {
# "name": "OrderItem",
# "type": "weak",
# "attributes": [
# {"name": "order_id", "type": "primary_key"},
# {"name": "product_id", "type": "primary_key"},
# {"name": "quantity", "type": "regular"},
# {"name": "unit_price", "type": "regular"},
# {"name": "discount", "type": "regular"}
# ]
# },
# {
# "name": "Payment",
# "type": "strong",
# "attributes": [
# {"name": "payment_id", "type": "primary_key"},
# {"name": "order_id", "type": "foreign_key"},
# {"name": "payment_method", "type": "regular"},
# {"name": "amount", "type": "regular"},
# {"name": "payment_date", "type": "regular"},
# {"name": "transaction_id", "type": "regular"}
# ]
# },
# {
# "name": "Supplier",
# "type": "strong",
# "attributes": [
# {"name": "supplier_id", "type": "primary_key"},
# {"name": "company_name", "type": "regular"},
# {"name": "contact_person", "type": "regular"},
# {"name": "email", "type": "regular"},
# {"name": "phone", "type": "multivalued", "multivalued": true},
# {"name": "address", "type": "composite", "composite": true}
# ]
# },
# {
# "name": "Review",
# "type": "weak",
# "attributes": [
# {"name": "customer_id", "type": "primary_key"},
# {"name": "product_id", "type": "primary_key"},
# {"name": "rating", "type": "regular"},
# {"name": "comment", "type": "regular"},
# {"name": "review_date", "type": "regular"}
# ]
# }
# ],
# "relationships": [
# {
# "name": "Places",
# "type": "regular",
# "entities": ["Customer", "Order"],
# "cardinalities": {
# "Customer": "1",
# "Order": "N"
# }
# },
# {
# "name": "Contains",
# "type": "identifying",
# "entities": ["Order", "OrderItem"],
# "cardinalities": {
# "Order": "1",
# "OrderItem": "N"
# }
# },
# {
# "name": "Includes",
# "type": "regular",
# "entities": ["OrderItem", "Product"],
# "cardinalities": {
# "OrderItem": "N",
# "Product": "1"
# }
# },
# {
# "name": "BelongsTo",
# "type": "regular",
# "entities": ["Product", "Category"],
# "cardinalities": {
# "Product": "N",
# "Category": "1"
# }
# },
# {
# "name": "ProcessedBy",
# "type": "regular",
# "entities": ["Order", "Payment"],
# "cardinalities": {
# "Order": "1",
# "Payment": "1"
# }
# },
# {
# "name": "Supplies",
# "type": "regular",
# "entities": ["Supplier", "Product"],
# "cardinalities": {
# "Supplier": "N",
# "Product": "M"
# }
# },
# {
# "name": "Writes",
# "type": "weak",
# "entities": ["Customer", "Review"],
# "cardinalities": {
# "Customer": "1",
# "Review": "N"
# }
# },
# {
# "name": "ReviewsProduct",
# "type": "weak",
# "entities": ["Review", "Product"],
# "cardinalities": {
# "Review": "N",
# "Product": "1"
# }
# },
# {
# "name": "Subcategory",
# "type": "regular",
# "entities": ["Category", "Category"],
# "cardinalities": {
# "Category": "1",
# "Category": "N"
# }
# }
# ]
# }
# Returns:
# str: The filepath to the generated PNG image file.
# """
# 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='dot')
# dot.attr(rankdir='TB', bgcolor='white', pad='0.5')
# dot.attr('node', fontname='Arial', fontsize='10', color='black')
# dot.attr('edge', fontname='Arial', fontsize='9', color='black')
# 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
# label_parts = [entity_name]
# if attributes:
# attr_lines = []
# for attr in attributes:
# attr_name = attr.get('name', '')
# 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:
# attr_display = f"PK: {{ {attr_name} }}"
# else:
# attr_display = f"PK: {attr_name}"
# elif attr_type == 'foreign_key':
# attr_display = f"FK: {attr_name}"
# else:
# attr_display = attr_name
# if is_derived:
# attr_display = f"/ {attr_display} /"
# if is_multivalued:
# attr_display = f"{{ {attr_display} }}"
# if is_composite:
# attr_display = f"( {attr_display} )"
# attr_lines.append(attr_display)
# if attr_lines:
# label_parts.extend(attr_lines)
# label = "\\n".join(label_parts)
# if entity_type == 'weak':
# style = 'filled'
# fillcolor = '#e8e8e8'
# penwidth = '2'
# else:
# style = 'filled'
# fillcolor = '#d0d0d0'
# penwidth = '1'
# dot.node(entity_name, label, shape='box', style=style, fillcolor=fillcolor, penwidth=penwidth)
# for relationship in relationships:
# rel_name = relationship.get('name')
# rel_type = relationship.get('type', 'regular')
# entities_involved = relationship.get('entities', [])
# cardinalities = relationship.get('cardinalities', {})
# if not rel_name or len(entities_involved) < 2:
# continue
# if rel_type == 'identifying':
# style = 'filled'
# fillcolor = '#c0c0c0'
# penwidth = '2'
# else:
# style = 'filled'
# fillcolor = '#c0c0c0'
# penwidth = '1'
# dot.node(rel_name, rel_name, shape='diamond', style=style, fillcolor=fillcolor, penwidth=penwidth)
# for entity in entities_involved:
# cardinality = cardinalities.get(entity, '1')
# dot.edge(entity, rel_name, label=cardinality)
# 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)}"
import graphviz
import json
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"
}
],
"relationships": [
{
"name": "Places",
"type": "regular",
"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"
data = json.loads(json_input)
if 'entities' not in data:
raise ValueError("Missing required field: entities")
# Configuración simple y funcional
dot = graphviz.Graph(comment='ER Diagram', engine='neato')
dot.attr(
bgcolor='white',
pad='1.0',
overlap='false',
splines='true',
sep='+15'
)
dot.attr('node', fontname='Arial', fontsize='11', color='#404040')
dot.attr('edge', fontname='Arial', fontsize='10', color='#404040')
entities = data.get('entities', [])
relationships = data.get('relationships', [])
# Crear solo las entidades principales (sin atributos por ahora)
for entity in entities:
entity_name = entity.get('name')
entity_type = entity.get('type', 'strong')
if not entity_name:
continue
if entity_type == 'weak':
# Entidad débil: rectángulo con doble borde
dot.node(
entity_name,
entity_name,
shape='box',
style='filled',
fillcolor='#e8e8e8',
color='#404040',
penwidth='3',
width='1.5',
height='0.8'
)
else:
# Entidad fuerte: rectángulo simple
dot.node(
entity_name,
entity_name,
shape='box',
style='filled',
fillcolor='#e8e8e8',
color='#404040',
penwidth='1',
width='1.5',
height='0.8'
)
# Crear relaciones
for relationship in relationships:
rel_name = relationship.get('name')
rel_type = relationship.get('type', 'regular')
entities_involved = relationship.get('entities', [])
cardinalities = relationship.get('cardinalities', {})
if not rel_name:
continue
if rel_type == 'isa':
# Relación ISA - crear triángulo y conectar correctamente
parent = relationship.get('parent')
children = relationship.get('children', [])
if parent and children:
# Crear triángulo ISA
isa_node = f"ISA_{rel_name}"
dot.node(
isa_node,
'ISA',
shape='triangle',
style='filled',
fillcolor='#d0d0d0',
color='#404040',
penwidth='2',
width='0.8',
height='0.6'
)
# Conectar padre con triángulo ISA
dot.edge(parent, isa_node, color='#404040', penwidth='2')
# Conectar triángulo ISA con cada hijo
for child in children:
dot.edge(isa_node, child, color='#404040', penwidth='2')
elif len(entities_involved) >= 2:
# Relaciones regulares e identificadoras
entity1, entity2 = entities_involved[0], entities_involved[1]
card1 = cardinalities.get(entity1, '1')
card2 = cardinalities.get(entity2, '1')
if rel_type == 'identifying':
# Relación identificadora - rombo con doble borde
dot.node(
rel_name,
rel_name,
shape='diamond',
style='filled',
fillcolor='#d8d8d8',
color='#404040',
penwidth='3',
width='1.5',
height='0.8'
)
else:
# Relación regular - rombo simple
dot.node(
rel_name,
rel_name,
shape='diamond',
style='filled',
fillcolor='#d8d8d8',
color='#404040',
penwidth='1',
width='1.5',
height='0.8'
)
# Conectar primera entidad con relación
dot.edge(
entity1,
rel_name,
taillabel=f' {card1} ',
color='#404040',
penwidth='2',
fontcolor='#000000',
fontsize='12'
)
# Conectar relación con segunda entidad
dot.edge(
rel_name,
entity2,
headlabel=f' {card2} ',
color='#404040',
penwidth='2',
fontcolor='#000000',
fontsize='12'
)
# 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}"
except json.JSONDecodeError:
return "Error: Invalid JSON format"
except Exception as e:
return f"Error: {str(e)}"