Spaces:
Running
Running
| 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)}" |