Update entity_relationship_generator.py

entity_relationship_generator.py

@@ -4,6 +4,385 @@ from tempfile import NamedTemporaryFile
 import os
 
 def generate_entity_relationship_diagram(json_input: str, output_format: str) -> str:
+    """
+    Generates an Entity Relationship (ER) diagram from JSON input.
+
+    Args:
+        json_input (str): A JSON string describing the ER diagram structure.
+                          It must follow the Expected JSON Format Example below.
+
+    Expected JSON Format Example:
+    {
+      "entities": [
+        {
+          "name": "User",
+          "type": "strong",
+          "attributes": [
+            {
+              "name": "user_id",
+              "type": "primary_key"
+            },
+            {
+              "name": "username",
+              "type": "regular"
+            },
+            {
+              "name": "email",
+              "type": "regular"
+            },
+            {
+              "name": "password_hash",
+              "type": "regular"
+            },
+            {
+              "name": "full_name",
+              "type": "composite"
+            },
+            {
+              "name": "phone_numbers",
+              "type": "multivalued"
+            },
+            {
+              "name": "age",
+              "type": "derived"
+            }
+          ]
+        },
+        {
+          "name": "Product",
+          "type": "strong",
+          "attributes": [
+            {
+              "name": "product_id",
+              "type": "primary_key"
+            },
+            {
+              "name": "name",
+              "type": "regular"
+            },
+            {
+              "name": "description",
+              "type": "regular"
+            },
+            {
+              "name": "price",
+              "type": "regular"
+            },
+            {
+              "name": "stock_quantity",
+              "type": "regular"
+            },
+            {
+              "name": "tags",
+              "type": "multivalued"
+            }
+          ]
+        },
+        {
+          "name": "Category",
+          "type": "strong",
+          "attributes": [
+            {
+              "name": "category_id",
+              "type": "primary_key"
+            },
+            {
+              "name": "name",
+              "type": "regular"
+            },
+            {
+              "name": "description",
+              "type": "regular"
+            }
+          ]
+        },
+        {
+          "name": "Order",
+          "type": "strong",
+          "attributes": [
+            {
+              "name": "order_id",
+              "type": "primary_key"
+            },
+            {
+              "name": "order_date",
+              "type": "regular"
+            },
+            {
+              "name": "status",
+              "type": "regular"
+            },
+            {
+              "name": "total_amount",
+              "type": "derived"
+            },
+            {
+              "name": "shipping_address",
+              "type": "composite"
+            }
+          ]
+        },
+        {
+          "name": "OrderItem",
+          "type": "weak",
+          "attributes": [
+            {
+              "name": "line_number",
+              "type": "partial_key"
+            },
+            {
+              "name": "quantity",
+              "type": "regular"
+            },
+            {
+              "name": "unit_price",
+              "type": "regular"
+            },
+            {
+              "name": "subtotal",
+              "type": "derived"
+            }
+          ]
+        },
+        {
+          "name": "Payment",
+          "type": "strong",
+          "attributes": [
+            {
+              "name": "payment_id",
+              "type": "primary_key"
+            },
+            {
+              "name": "amount",
+              "type": "regular"
+            },
+            {
+              "name": "payment_method",
+              "type": "regular"
+            },
+            {
+              "name": "payment_date",
+              "type": "regular"
+            },
+            {
+              "name": "status",
+              "type": "regular"
+            }
+          ]
+        },
+        {
+          "name": "Review",
+          "type": "strong",
+          "attributes": [
+            {
+              "name": "review_id",
+              "type": "primary_key"
+            },
+            {
+              "name": "rating",
+              "type": "regular"
+            },
+            {
+              "name": "comment",
+              "type": "regular"
+            },
+            {
+              "name": "review_date",
+              "type": "regular"
+            }
+          ]
+        },
+        {
+          "name": "Vendor",
+          "type": "strong",
+          "attributes": [
+            {
+              "name": "vendor_id",
+              "type": "primary_key"
+            },
+            {
+              "name": "company_name",
+              "type": "regular"
+            },
+            {
+              "name": "contact_person",
+              "type": "regular"
+            },
+            {
+              "name": "contact_emails",
+              "type": "multivalued"
+            },
+            {
+              "name": "business_address",
+              "type": "composite"
+            }
+          ]
+        },
+        {
+          "name": "ShoppingCart",
+          "type": "strong",
+          "attributes": [
+            {
+              "name": "cart_id",
+              "type": "primary_key"
+            },
+            {
+              "name": "created_date",
+              "type": "regular"
+            },
+            {
+              "name": "last_updated",
+              "type": "regular"
+            },
+            {
+              "name": "total_items",
+              "type": "derived"
+            }
+          ]
+        },
+        {
+          "name": "CartItem",
+          "type": "weak",
+          "attributes": [
+            {
+              "name": "item_position",
+              "type": "partial_key"
+            },
+            {
+              "name": "quantity",
+              "type": "regular"
+            },
+            {
+              "name": "added_date",
+              "type": "regular"
+            }
+          ]
+        }
+      ],
+      "relationships": [
+        {
+          "name": "PlacesOrder",
+          "type": "regular",
+          "entities": ["User", "Order"],
+          "cardinalities": {
+            "User": "1",
+            "Order": "M"
+          },
+          "attributes": []
+        },
+        {
+          "name": "Contains",
+          "type": "identifying",
+          "entities": ["Order", "OrderItem"],
+          "cardinalities": {
+            "Order": "1",
+            "OrderItem": "M"
+          },
+          "attributes": []
+        },
+        {
+          "name": "OrdersProduct",
+          "type": "regular",
+          "entities": ["OrderItem", "Product"],
+          "cardinalities": {
+            "OrderItem": "M",
+            "Product": "1"
+          },
+          "attributes": []
+        },
+        {
+          "name": "BelongsTo",
+          "type": "regular",
+          "entities": ["Product", "Category"],
+          "cardinalities": {
+            "Product": "M",
+            "Category": "1"
+          },
+          "attributes": []
+        },
+        {
+          "name": "ProcessesPayment",
+          "type": "regular",
+          "entities": ["Order", "Payment"],
+          "cardinalities": {
+            "Order": "1",
+            "Payment": "M"
+          },
+          "attributes": []
+        },
+        {
+          "name": "WritesReview",
+          "type": "regular",
+          "entities": ["User", "Review"],
+          "cardinalities": {
+            "User": "1",
+            "Review": "M"
+          },
+          "attributes": []
+        },
+        {
+          "name": "ReviewsProduct",
+          "type": "regular",
+          "entities": ["Review", "Product"],
+          "cardinalities": {
+            "Review": "M",
+            "Product": "1"
+          },
+          "attributes": []
+        },
+        {
+          "name": "Supplies",
+          "type": "regular",
+          "entities": ["Vendor", "Product"],
+          "cardinalities": {
+            "Vendor": "M",
+            "Product": "M"
+          },
+          "attributes": [
+            {
+              "name": "supply_price"
+            },
+            {
+              "name": "lead_time"
+            }
+          ]
+        },
+        {
+          "name": "HasCart",
+          "type": "regular",
+          "entities": ["User", "ShoppingCart"],
+          "cardinalities": {
+            "User": "1",
+            "ShoppingCart": "1"
+          },
+          "attributes": []
+        },
+        {
+          "name": "CartContains",
+          "type": "identifying",
+          "entities": ["ShoppingCart", "CartItem"],
+          "cardinalities": {
+            "ShoppingCart": "1",
+            "CartItem": "M"
+          },
+          "attributes": []
+        },
+        {
+          "name": "CartHasProduct",
+          "type": "regular",
+          "entities": ["CartItem", "Product"],
+          "cardinalities": {
+            "CartItem": "M",
+            "Product": "1"
+          },
+          "attributes": []
+        }
+      ]
+    }
+
+    Returns:
+        str: The filepath to the generated PNG image file.
+    """
     try:
         if not json_input.strip():
             return "Error: Empty input"
@@ -25,33 +404,12 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
         dot.attr('node', fontname='Arial', fontsize='10', color='#404040')
         dot.attr('edge', fontname='Arial', fontsize='9', color='#4a4a4a')
         
-        # Base color system - much lighter grays starting from #BEBEBE
-        base_color = '#BEBEBE'  # Much lighter base color
-        lightening_factor = 0.08  # Smaller factor for subtle gradations
+        entity_color = '#BEBEBE'        
+        attribute_color = '#D4D4D4'     
+        relationship_color = '#E8E8E8'  
+        isa_color = '#F0F0F0'          
         
-        def get_gradient_color(depth, base_hex_color, lightening_factor=0.08):
-            """Get lightened color based on depth - much lighter grays"""
-            if not isinstance(base_hex_color, str) or not base_hex_color.startswith('#') or len(base_hex_color) != 7:
-                base_hex_color = '#BEBEBE'  # Light gray fallback
-                
-            base_r = int(base_hex_color[1:3], 16)
-            base_g = int(base_hex_color[3:5], 16)
-            base_b = int(base_hex_color[5:7], 16)
-            
-            current_r = base_r + int((255 - base_r) * depth * lightening_factor)
-            current_g = base_g + int((255 - base_g) * depth * lightening_factor)
-            current_b = base_b + int((255 - base_b) * depth * lightening_factor)
-            
-            current_r = min(255, current_r)
-            current_g = min(255, current_g)
-            current_b = min(255, current_b)
-            
-            return f'#{current_r:02x}{current_g:02x}{current_b:02x}'
-
-        def get_font_color_for_background(depth, lightening_factor=0.08):
-            """Get appropriate font color based on background lightness"""
-            # With lighter base colors, we'll use black text for better readability
-            return 'black'
+        font_color = 'black'
         
         entities = data.get('entities', [])
         relationships = data.get('relationships', [])
@@ -65,18 +423,14 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             if not entity_name:
                 continue
             
-            # Entity colors - depth 1 for entities
-            entity_color = get_gradient_color(1, base_color, lightening_factor)
-            entity_font_color = get_font_color_for_background(1, lightening_factor)
-            
             if entity_type == 'weak':
                 dot.node(
                     entity_name, 
                     entity_name, 
                     shape='box', 
-                    style='filled,rounded',  # KEY CHANGE: rounded borders like your other generators
-                    fillcolor=entity_color,
-                    fontcolor=entity_font_color,
+                    style='filled,rounded',
+                    fillcolor=entity_color,  
+                    fontcolor=font_color,
                     color='#404040',
                     penwidth='3',
                     width='1.8',
@@ -88,9 +442,9 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                     entity_name, 
                     entity_name, 
                     shape='box', 
-                    style='filled,rounded',  # KEY CHANGE: rounded borders
-                    fillcolor=entity_color,
-                    fontcolor=entity_font_color,
+                    style='filled,rounded',
+                    fillcolor=entity_color,  
+                    fontcolor=font_color,
                     color='#404040',
                     penwidth='1',
                     width='1.8',
@@ -98,23 +452,20 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                     fontsize='12'
                 )
             
-            # Process attributes with gradient colors - depth 2 for attributes
             for i, attr in enumerate(attributes):
                 attr_name = attr.get('name', '')
                 attr_type = attr.get('type', 'regular')
                 
                 attr_id = f"{entity_name}_attr_{i}"
-                attr_color = get_gradient_color(2, base_color, lightening_factor)
-                attr_font_color = get_font_color_for_background(2, lightening_factor)
                 
                 if attr_type == 'primary_key':
                     dot.node(
                         attr_id, 
                         f'{attr_name} (PK)', 
                         shape='ellipse', 
-                        style='filled,rounded',  # KEY CHANGE: rounded borders
-                        fillcolor=attr_color,
-                        fontcolor=attr_font_color,
+                        style='filled,rounded',
+                        fillcolor=attribute_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         width='1.2',
                         height='0.6',
@@ -125,9 +476,9 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         attr_id, 
                         f'{attr_name} (Partial)', 
                         shape='ellipse', 
-                        style='filled,rounded,dashed',  # KEY CHANGE: rounded borders
-                        fillcolor=attr_color,
-                        fontcolor=attr_font_color,
+                        style='filled,rounded,dashed',
+                        fillcolor=attribute_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         width='1.2',
                         height='0.6',
@@ -138,9 +489,9 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         attr_id, 
                         attr_name, 
                         shape='ellipse', 
-                        style='filled,rounded',  # KEY CHANGE: rounded borders
-                        fillcolor=attr_color,
-                        fontcolor=attr_font_color,
+                        style='filled,rounded',
+                        fillcolor=attribute_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         penwidth='3',
                         width='1.2',
@@ -152,9 +503,9 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         attr_id, 
                         f'/{attr_name}/', 
                         shape='ellipse', 
-                        style='filled,rounded,dashed',  # KEY CHANGE: rounded borders
-                        fillcolor=attr_color,
-                        fontcolor=attr_font_color,
+                        style='filled,rounded,dashed',
+                        fillcolor=attribute_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         width='1.2',
                         height='0.6',
@@ -165,22 +516,22 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         attr_id, 
                         attr_name, 
                         shape='ellipse', 
-                        style='filled,rounded',  # KEY CHANGE: rounded borders
-                        fillcolor=attr_color,
-                        fontcolor=attr_font_color,
+                        style='filled,rounded',
+                        fillcolor=attribute_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         width='1.2',
                         height='0.6',
                         fontsize='10'
                     )
-                else:  # regular
+                else:  
                     dot.node(
                         attr_id, 
                         attr_name, 
                         shape='ellipse', 
-                        style='filled,rounded',  # KEY CHANGE: rounded borders
-                        fillcolor=attr_color,
-                        fontcolor=attr_font_color,
+                        style='filled,rounded',
+                        fillcolor=attribute_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         width='1.2',
                         height='0.6',
@@ -189,7 +540,6 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                 
                 dot.edge(entity_name, attr_id, color='#4a4a4a', len='1.5')
         
-        # Process relationships with gradient colors - depth 1.5 for relationships
         for relationship in relationships:
             rel_name = relationship.get('name')
             rel_type = relationship.get('type', 'regular')
@@ -200,25 +550,20 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
             if not rel_name:
                 continue
             
-            rel_color = get_gradient_color(1.5, base_color, lightening_factor)
-            rel_font_color = get_font_color_for_background(1.5, lightening_factor)
-            
             if rel_type == 'isa':
                 parent = relationship.get('parent')
                 children = relationship.get('children', [])
                 
                 if parent and children:
                     isa_id = f"isa_{rel_name}"
-                    isa_color = get_gradient_color(1.2, base_color, lightening_factor)
-                    isa_font_color = get_font_color_for_background(1.2, lightening_factor)
                     
                     dot.node(
                         isa_id, 
                         'ISA', 
                         shape='triangle', 
-                        style='filled,rounded',  # KEY CHANGE: rounded borders
-                        fillcolor=isa_color,
-                        fontcolor=isa_font_color,
+                        style='filled,rounded',
+                        fillcolor=isa_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         penwidth='2',
                         width='1.0',
@@ -237,9 +582,9 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         rel_name, 
                         rel_name, 
                         shape='diamond', 
-                        style='filled,rounded',  # KEY CHANGE: rounded borders
-                        fillcolor=rel_color,
-                        fontcolor=rel_font_color,
+                        style='filled,rounded',
+                        fillcolor=relationship_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         penwidth='3',
                         width='1.8',
@@ -251,9 +596,9 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         rel_name, 
                         rel_name, 
                         shape='diamond', 
-                        style='filled,rounded',  # KEY CHANGE: rounded borders
-                        fillcolor=rel_color,
-                        fontcolor=rel_font_color,
+                        style='filled,rounded',
+                        fillcolor=relationship_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         penwidth='1',
                         width='1.8',
@@ -261,20 +606,17 @@ def generate_entity_relationship_diagram(json_input: str, output_format: str) ->
                         fontsize='11'
                     )
                 
-                # Relationship attributes - depth 2.5
                 for j, attr in enumerate(rel_attributes):
                     attr_name = attr.get('name', '')
                     attr_id = f"{rel_name}_attr_{j}"
-                    rel_attr_color = get_gradient_color(2.5, base_color, lightening_factor)
-                    rel_attr_font_color = get_font_color_for_background(2.5, lightening_factor)
                     
                     dot.node(
                         attr_id, 
                         attr_name, 
                         shape='ellipse', 
-                        style='filled,rounded',  # KEY CHANGE: rounded borders
-                        fillcolor=rel_attr_color,
-                        fontcolor=rel_attr_font_color,
+                        style='filled,rounded',
+                        fillcolor=attribute_color,  
+                        fontcolor=font_color,
                         color='#404040',
                         width='1.0',
                         height='0.5',