Create entity_relationship_generator.py

entity_relationship_generator.py

@@ -0,0 +1,181 @@
+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(
+            name='ERDiagram',
+            format='png',
+            graph_attr={
+                'rankdir': 'TB',
+                'splines': 'ortho',
+                'bgcolor': 'white',
+                'pad': '0.5',
+                'nodesep': '2.0',
+                'ranksep': '2.5'
+            }
+        )
+        
+        base_color = '#19191a'
+        lightening_factor = 0.15
+        
+        entities = data.get('entities', [])
+        relationships = data.get('relationships', [])
+        
+        for i, entity in enumerate(entities):
+            entity_name = entity.get('name')
+            entity_type = entity.get('type', 'strong')
+            attributes = entity.get('attributes', [])
+            
+            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"
+            
+            if attributes:
+                primary_keys = []
+                foreign_keys = []
+                regular_attrs = []
+                
+                for attr in attributes:
+                    attr_name = attr.get('name', '')
+                    attr_type = attr.get('type', 'key')
+                    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} }}}}")
+                        else:
+                            primary_keys.append(f"[PK] {attr_name}")
+                    elif attr_type == 'foreign_key':
+                        foreign_keys.append(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} )"
+                        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)
+            
+            if entity_type == 'weak':
+                shape = 'doubleoctagon'
+                style = 'filled'
+            else:
+                shape = 'box'
+                style = 'filled,rounded'
+            
+            dot.node(
+                entity_name,
+                entity_label,
+                shape=shape,
+                style=style,
+                fillcolor=node_color,
+                fontcolor=font_color,
+                fontsize='10',
+                fontname='Helvetica'
+            )
+        
+        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:
+                raise ValueError(f"Invalid relationship: {relationship}")
+            
+            rel_node_color = '#e6f3ff'
+            
+            if rel_type == 'identifying':
+                rel_shape = 'diamond'
+                rel_style = 'filled,bold'
+                rel_color = '#4a90e2'
+            elif rel_type == 'weak':
+                rel_shape = 'diamond'
+                rel_style = 'filled,dashed'
+                rel_color = '#a0a0a0'
+            else:
+                rel_shape = 'diamond'
+                rel_style = 'filled'
+                rel_color = '#4a90e2'
+            
+            dot.node(
+                rel_name,
+                rel_name,
+                shape=rel_shape,
+                style=rel_style,
+                fillcolor=rel_color,
+                fontcolor='white',
+                fontsize='10',
+                fontname='Helvetica'
+            )
+            
+            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'
+                )
+
+        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)}"