Create class_diagram_generator.py

class_diagram_generator.py

@@ -0,0 +1,171 @@
+import graphviz
+import json
+from tempfile import NamedTemporaryFile
+import os
+
+def generate_class_diagram(json_input: str, output_format: str) -> str:
+    try:
+        if not json_input.strip():
+            return "Error: Empty input"
+            
+        data = json.loads(json_input)
+        
+        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.5',
+                'ranksep': '2.0'
+            }
+        )
+        
+        base_color = '#19191a'
+        lightening_factor = 0.15
+        
+        classes = data.get('classes', [])
+        relationships = data.get('relationships', [])
+        
+        for i, cls in enumerate(classes):
+            class_name = cls.get('name')
+            class_type = cls.get('type', 'class')
+            attributes = cls.get('attributes', [])
+            methods = cls.get('methods', [])
+            
+            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':
+                class_label += "<<abstract>>\\n"
+            elif class_type == 'interface':
+                class_label += "<<interface>>\\n"
+            elif class_type == 'enum':
+                class_label += "<<enumeration>>\\n"
+            
+            class_label += f"{class_name}\\l"
+            
+            if attributes:
+                class_label += "\\l"
+                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}"
+                    if attr_type:
+                        attr_line += f" : {attr_type}"
+                    class_label += f"{attr_line}\\l"
+            
+            if methods:
+                class_label += "\\l"
+                for method in methods:
+                    visibility = method.get('visibility', '+')
+                    name = method.get('name', '')
+                    parameters = method.get('parameters', [])
+                    return_type = method.get('return_type', 'void')
+                    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}("
+                    if parameters:
+                        param_strs = []
+                        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}"
+                    class_label += f"{method_line}\\l"
+            
+            if class_type == 'interface':
+                style = 'filled,dashed'
+            else:
+                style = 'filled'
+            
+            dot.node(
+                class_name,
+                class_label,
+                shape='record',
+                style=style,
+                fillcolor=node_color,
+                fontcolor=font_color,
+                fontsize='10',
+                fontname='Helvetica'
+            )
+        
+        for relationship in relationships:
+            from_class = relationship.get('from')
+            to_class = relationship.get('to')
+            rel_type = relationship.get('type', 'association')
+            label = relationship.get('label', '')
+            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}")
+            
+            edge_label = label
+            if multiplicity_from or multiplicity_to:
+                edge_label += f"\\n{multiplicity_from} --- {multiplicity_to}"
+            
+            if rel_type == 'inheritance':
+                dot.edge(from_class, to_class, arrowhead='empty', color='#4a4a4a', label=edge_label, fontsize='9')
+            elif rel_type == 'composition':
+                dot.edge(from_class, to_class, arrowhead='normal', arrowtail='diamond', dir='both', color='#4a4a4a', label=edge_label, fontsize='9')
+            elif rel_type == 'aggregation':
+                dot.edge(from_class, to_class, arrowhead='normal', arrowtail='odiamond', dir='both', color='#4a4a4a', label=edge_label, fontsize='9')
+            elif rel_type == 'realization':
+                dot.edge(from_class, to_class, arrowhead='empty', style='dashed', color='#4a4a4a', label=edge_label, fontsize='9')
+            elif rel_type == 'dependency':
+                dot.edge(from_class, to_class, arrowhead='normal', style='dashed', color='#4a4a4a', label=edge_label, fontsize='9')
+            else:
+                dot.edge(from_class, to_class, arrowhead='normal', color='#4a4a4a', label=edge_label, fontsize='9')
+
+        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)}"