Update class_diagram_generator.py

class_diagram_generator.py

@@ -194,7 +194,10 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
         dot = graphviz.Digraph(comment='Class Diagram')
         dot.attr(rankdir='TB', bgcolor='white', pad='0.5', nodesep='0.8', ranksep='1.2', splines='ortho')
         dot.attr('node', shape='plaintext', fontname='Arial', fontsize='11')
-        dot.attr('edge', color='black', fontname='Arial', fontsize='9', minlen='1')
+        dot.attr('edge', color='#4a4a4a', fontname='Arial', fontsize='9', minlen='1')
+        
+        header_color = '#BEBEBE'
+        body_color = '#E8E8E8'
         
         classes = data.get('classes', [])
         relationships = data.get('relationships', [])
@@ -208,18 +211,16 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
             if not class_name:
                 continue
             
-            html_label = '<TABLE BORDER="1" CELLBORDER="0" CELLSPACING="0" CELLPADDING="5" BGCOLOR="white">'
+            html_label = f'<TABLE BORDER="1" CELLBORDER="0" CELLSPACING="0" CELLPADDING="5" STYLE="rounded" BGCOLOR="{body_color}">'
             
             if class_type == 'abstract':
-                html_label += f'<TR><TD ALIGN="CENTER"><B>&lt;&lt;abstract&gt;&gt;<BR/>{class_name}</B></TD></TR>'
+                html_label += f'<TR><TD ALIGN="CENTER" BGCOLOR="{header_color}"><B>&lt;&lt;abstract&gt;&gt;<BR/>{class_name}</B></TD></TR>'
             elif class_type == 'interface':
-                html_label += f'<TR><TD ALIGN="CENTER"><B>&lt;&lt;interface&gt;&gt;<BR/>{class_name}</B></TD></TR>'
+                html_label += f'<TR><TD ALIGN="CENTER" BGCOLOR="{header_color}"><B>&lt;&lt;interface&gt;&gt;<BR/>{class_name}</B></TD></TR>'
             elif class_type == 'enum':
-                html_label += f'<TR><TD ALIGN="CENTER"><B>&lt;&lt;enumeration&gt;&gt;<BR/>{class_name}</B></TD></TR>'
+                html_label += f'<TR><TD ALIGN="CENTER" BGCOLOR="{header_color}"><B>&lt;&lt;enumeration&gt;&gt;<BR/>{class_name}</B></TD></TR>'
             else:
-                html_label += f'<TR><TD ALIGN="CENTER"><B>{class_name}</B></TD></TR>'
-            
-            html_label += '<HR/>'
+                html_label += f'<TR><TD ALIGN="CENTER" BGCOLOR="{header_color}"><B>{class_name}</B></TD></TR>'
             
             if attributes:
                 for attr in attributes:
@@ -234,11 +235,9 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
                     if is_static:
                         line = f"<U>{line}</U>"
                     
-                    html_label += f'<TR><TD ALIGN="LEFT">{line}</TD></TR>'
+                    html_label += f'<TR><TD ALIGN="LEFT" BGCOLOR="{body_color}">{line}</TD></TR>'
             else:
-                html_label += '<TR><TD ALIGN="LEFT"> </TD></TR>'
-            
-            html_label += '<HR/>'
+                html_label += f'<TR><TD ALIGN="LEFT" BGCOLOR="{body_color}"> </TD></TR>'
             
             if methods:
                 for method in methods:
@@ -264,13 +263,13 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
                     if is_abstract:
                         line = f"<I>{line}</I>"
                     
-                    html_label += f'<TR><TD ALIGN="LEFT">{line}</TD></TR>'
+                    html_label += f'<TR><TD ALIGN="LEFT" BGCOLOR="{body_color}">{line}</TD></TR>'
             else:
-                html_label += '<TR><TD ALIGN="LEFT"> </TD></TR>'
+                html_label += f'<TR><TD ALIGN="LEFT" BGCOLOR="{body_color}"> </TD></TR>'
             
             html_label += '</TABLE>'
             
-            dot.node(class_name, f'<{html_label}>')
+            dot.node(class_name, f'<{html_label}>', style='rounded')
         
         for relationship in relationships:
             from_class = relationship.get('from')
@@ -284,7 +283,7 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
                 continue
             
             edge_attrs = {
-                'color': 'black'
+                'color': '#4a4a4a'
             }
             
             if multiplicity_from:
@@ -295,28 +294,28 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
             
             if rel_type == 'inheritance':
                 edge_attrs['arrowhead'] = 'empty'
-                edge_attrs['color'] = 'black'
+                edge_attrs['color'] = '#4a4a4a'
             elif rel_type == 'composition':
                 edge_attrs['arrowhead'] = 'normal'
                 edge_attrs['arrowtail'] = 'diamond'
                 edge_attrs['dir'] = 'both'
-                edge_attrs['color'] = 'black'
+                edge_attrs['color'] = '#4a4a4a'
             elif rel_type == 'aggregation':
                 edge_attrs['arrowhead'] = 'normal'
                 edge_attrs['arrowtail'] = 'odiamond'
                 edge_attrs['dir'] = 'both'
-                edge_attrs['color'] = 'black'
+                edge_attrs['color'] = '#4a4a4a'
             elif rel_type == 'realization':
                 edge_attrs['arrowhead'] = 'empty'
                 edge_attrs['style'] = 'dashed'
-                edge_attrs['color'] = 'black'
+                edge_attrs['color'] = '#4a4a4a'
             elif rel_type == 'dependency':
                 edge_attrs['arrowhead'] = 'normal'
                 edge_attrs['style'] = 'dashed'
-                edge_attrs['color'] = 'black'
+                edge_attrs['color'] = '#4a4a4a'
             else:
                 edge_attrs['arrowhead'] = 'normal'
-                edge_attrs['color'] = 'black'
+                edge_attrs['color'] = '#4a4a4a'
             
             dot.edge(from_class, to_class, **edge_attrs)