Update class_diagram_generator.py

class_diagram_generator.py

@@ -199,18 +199,10 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
         if 'classes' not in data:
             raise ValueError("Missing required field: classes")
 
-        dot = graphviz.Digraph(
-            name='ClassDiagram',
-            format='png',
-            graph_attr={
-                'rankdir': 'TB',
-                'splines': 'ortho',
-                'bgcolor': 'white',
-                'pad': '0.5',
-                'nodesep': '1.0',
-                'ranksep': '1.5'
-            }
-        )
+        dot = graphviz.Digraph(comment='Class Diagram')
+        dot.attr(rankdir='TB', bgcolor='white', pad='0.5')
+        dot.attr('node', shape='box', style='filled', fillcolor='white', color='black', fontname='Arial', fontsize='10')
+        dot.attr('edge', color='black', fontname='Arial', fontsize='9')
         
         classes = data.get('classes', [])
         relationships = data.get('relationships', [])
@@ -222,44 +214,40 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
             methods = cls.get('methods', [])
             
             if not class_name:
-                raise ValueError(f"Invalid class: {cls}")
+                continue
             
-            class_parts = []
+            label_parts = []
             
-            header = class_name
             if class_type == 'abstract':
-                header = f"<<abstract>>\\n{class_name}"
+                label_parts.append(f"&lt;&lt;abstract&gt;&gt;\\n{class_name}")
             elif class_type == 'interface':
-                header = f"<<interface>>\\n{class_name}"
+                label_parts.append(f"&lt;&lt;interface&gt;&gt;\\n{class_name}")
             elif class_type == 'enum':
-                header = f"<<enumeration>>\\n{class_name}"
-            
-            class_parts.append(header)
+                label_parts.append(f"&lt;&lt;enumeration&gt;&gt;\\n{class_name}")
+            else:
+                label_parts.append(class_name)
             
             if attributes:
-                attr_section = ""
+                attr_lines = []
                 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} {name}"
+                    line = f"{visibility} {name}"
                     if attr_type:
-                        attr_line += f" : {attr_type}"
+                        line += f" : {attr_type}"
                     if is_static:
-                        attr_line += " [static]"
+                        line += " [static]"
                     
-                    if attr_section:
-                        attr_section += "\\l"
-                    attr_section += attr_line
+                    attr_lines.append(line)
                 
-                if attr_section:
-                    attr_section += "\\l"
-                    class_parts.append(attr_section)
+                if attr_lines:
+                    label_parts.append("\\n".join(attr_lines))
             
             if methods:
-                method_section = ""
+                method_lines = []
                 for method in methods:
                     visibility = method.get('visibility', '+')
                     name = method.get('name', '')
@@ -268,49 +256,39 @@ 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} {name}("
+                    line = f"{visibility} {name}("
                     if parameters:
-                        param_strs = []
+                        params = []
                         for param in parameters:
                             param_name = param.get('name', '')
                             param_type = param.get('type', '')
-                            param_strs.append(f"{param_name}: {param_type}")
-                        method_line += ", ".join(param_strs)
-                    method_line += f") : {return_type}"
+                            params.append(f"{param_name}: {param_type}")
+                        line += ", ".join(params)
+                    line += f") : {return_type}"
                     
                     if is_static:
-                        method_line += " [static]"
+                        line += " [static]"
                     if is_abstract:
-                        method_line += " [abstract]"
+                        line += " [abstract]"
                     
-                    if method_section:
-                        method_section += "\\l"
-                    method_section += method_line
+                    method_lines.append(line)
                 
-                if method_section:
-                    method_section += "\\l"
-                    class_parts.append(method_section)
+                if method_lines:
+                    label_parts.append("\\n".join(method_lines))
             
-            class_label = "|".join(class_parts)
+            label = "\\n\\n".join(label_parts)
             
             if class_type == 'interface':
-                style = 'filled,dashed'
                 fillcolor = '#f5f5f5'
+                style = 'filled,dashed'
             elif class_type == 'abstract':
-                style = 'filled'
                 fillcolor = '#eeeeee'
-            else:
                 style = 'filled'
+            else:
                 fillcolor = 'white'
+                style = 'filled'
             
-            dot.node(
-                class_name,
-                class_label,
-                shape='record',
-                style=style,
-                fillcolor=fillcolor,
-                color='black'
-            )
+            dot.node(class_name, label, fillcolor=fillcolor, style=style)
         
         for relationship in relationships:
             from_class = relationship.get('from')
@@ -320,10 +298,10 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
             multiplicity_from = relationship.get('multiplicity_from', '')
             multiplicity_to = relationship.get('multiplicity_to', '')
             
-            if not all([from_class, to_class]):
-                raise ValueError(f"Invalid relationship: {relationship}")
+            if not from_class or not to_class:
+                continue
             
-            edge_attrs = {'color': 'black'}
+            edge_attrs = {}
             
             if label:
                 edge_attrs['label'] = label