Update class_diagram_generator.py

class_diagram_generator.py

@@ -207,18 +207,25 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
                 'splines': 'ortho',
                 'bgcolor': 'white',
                 'pad': '0.5',
-                'nodesep': '1.5',
-                'ranksep': '2.0'
+                'nodesep': '1.0',
+                'ranksep': '1.5',
+                'fontname': 'Arial',
+                'dpi': '300'
+            },
+            node_attr={
+                'fontname': 'Arial',
+                'fontsize': '10'
+            },
+            edge_attr={
+                'fontname': 'Arial',
+                'fontsize': '9'
             }
         )
         
-        base_color = '#19191a'
-        lightening_factor = 0.15
-        
         classes = data.get('classes', [])
         relationships = data.get('relationships', [])
         
-        for i, cls in enumerate(classes):
+        for cls in classes:
             class_name = cls.get('name')
             class_type = cls.get('type', 'class')
             attributes = cls.get('attributes', [])
@@ -227,28 +234,6 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
             if not class_name:
                 raise ValueError(f"Invalid class: {cls}")
             
-            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'
-            
             class_label = ""
             
             if class_type == 'abstract':
@@ -258,26 +243,28 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
             elif class_type == 'enum':
                 class_label += "<<enumeration>>\\n"
             
-            class_label += f"{class_name}\\l"
+            class_label += f"<B>{class_name}</B>"
             
             if attributes:
-                class_label += "\\l"
+                class_label += "|"
                 for attr in attributes:
                     visibility = attr.get('visibility', '+')
                     name = attr.get('name', '')
                     attr_type = attr.get('type', '')
                     is_static = attr.get('static', False)
                     
-                    attr_line = f"{visibility} "
-                    if is_static:
-                        attr_line += f"<<static>> "
-                    attr_line += f"{name}"
+                    attr_line = f"{visibility} {name}"
                     if attr_type:
                         attr_line += f" : {attr_type}"
+                    if is_static:
+                        attr_line = f"<U>{attr_line}</U>"
                     class_label += f"{attr_line}\\l"
             
             if methods:
-                class_label += "\\l"
+                if not attributes:
+                    class_label += "|"
+                else:
+                    class_label += "|"
                 for method in methods:
                     visibility = method.get('visibility', '+')
                     name = method.get('name', '')
@@ -286,13 +273,7 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
                     is_static = method.get('static', False)
                     is_abstract = method.get('abstract', False)
                     
-                    method_line = f"{visibility} "
-                    if is_static:
-                        method_line += f"<<static>> "
-                    if is_abstract:
-                        method_line += f"<<abstract>> "
-                    
-                    method_line += f"{name}("
+                    method_line = f"{visibility} {name}("
                     if parameters:
                         param_strs = []
                         for param in parameters:
@@ -301,22 +282,32 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
                             param_strs.append(f"{param_name}: {param_type}")
                         method_line += ", ".join(param_strs)
                     method_line += f") : {return_type}"
+                    
+                    if is_static:
+                        method_line = f"<U>{method_line}</U>"
+                    if is_abstract:
+                        method_line = f"<I>{method_line}</I>"
+                    
                     class_label += f"{method_line}\\l"
             
             if class_type == 'interface':
                 style = 'filled,dashed'
+                fillcolor = '#f0f0f0'
+            elif class_type == 'abstract':
+                style = 'filled'
+                fillcolor = '#e8e8e8'
             else:
                 style = 'filled'
+                fillcolor = 'white'
             
             dot.node(
                 class_name,
-                class_label,
+                f"<{class_label}>",
                 shape='record',
                 style=style,
-                fillcolor=node_color,
-                fontcolor=font_color,
-                fontsize='10',
-                fontname='Helvetica'
+                fillcolor=fillcolor,
+                color='black',
+                penwidth='1'
             )
         
         for relationship in relationships:
@@ -330,22 +321,45 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
             if not all([from_class, to_class]):
                 raise ValueError(f"Invalid relationship: {relationship}")
             
-            edge_label = label
-            if multiplicity_from or multiplicity_to:
-                edge_label += f"\\n{multiplicity_from} --- {multiplicity_to}"
+            edge_attrs = {
+                'color': 'black',
+                'fontcolor': 'black',
+                'fontsize': '9'
+            }
+            
+            if label:
+                edge_attrs['label'] = label
+            
+            if multiplicity_from and multiplicity_to:
+                edge_attrs['headlabel'] = multiplicity_to
+                edge_attrs['taillabel'] = multiplicity_from
             
             if rel_type == 'inheritance':
-                dot.edge(from_class, to_class, arrowhead='empty', color='#4a4a4a', label=edge_label, fontsize='9')
+                edge_attrs['arrowhead'] = 'empty'
+                edge_attrs['arrowsize'] = '1.2'
             elif rel_type == 'composition':
-                dot.edge(from_class, to_class, arrowhead='normal', arrowtail='diamond', dir='both', color='#4a4a4a', label=edge_label, fontsize='9')
+                edge_attrs['arrowhead'] = 'normal'
+                edge_attrs['arrowtail'] = 'diamond'
+                edge_attrs['dir'] = 'both'
+                edge_attrs['arrowsize'] = '1.0'
             elif rel_type == 'aggregation':
-                dot.edge(from_class, to_class, arrowhead='normal', arrowtail='odiamond', dir='both', color='#4a4a4a', label=edge_label, fontsize='9')
+                edge_attrs['arrowhead'] = 'normal'
+                edge_attrs['arrowtail'] = 'odiamond'
+                edge_attrs['dir'] = 'both'
+                edge_attrs['arrowsize'] = '1.0'
             elif rel_type == 'realization':
-                dot.edge(from_class, to_class, arrowhead='empty', style='dashed', color='#4a4a4a', label=edge_label, fontsize='9')
+                edge_attrs['arrowhead'] = 'empty'
+                edge_attrs['style'] = 'dashed'
+                edge_attrs['arrowsize'] = '1.2'
             elif rel_type == 'dependency':
-                dot.edge(from_class, to_class, arrowhead='normal', style='dashed', color='#4a4a4a', label=edge_label, fontsize='9')
+                edge_attrs['arrowhead'] = 'normal'
+                edge_attrs['style'] = 'dashed'
+                edge_attrs['arrowsize'] = '1.0'
             else:
-                dot.edge(from_class, to_class, arrowhead='normal', color='#4a4a4a', label=edge_label, fontsize='9')
+                edge_attrs['arrowhead'] = 'normal'
+                edge_attrs['arrowsize'] = '1.0'
+            
+            dot.edge(from_class, to_class, **edge_attrs)
 
         with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
             dot.render(tmp.name, format=output_format, cleanup=True)