class_diagram_generator.py

import graphviz
import json
from tempfile import NamedTemporaryFile
import os

def generate_class_diagram(json_input: str, output_format: str) -> str:
    """
    Generates a class diagram from JSON input.

    Args:
        json_input (str): A JSON string describing the class diagram structure.
                          It must follow the Expected JSON Format Example below.

    Expected JSON Format Example:
    {
      "classes": [
        {
          "name": "Vehicle",
          "type": "abstract",
          "attributes": [
            {"name": "id", "type": "String", "visibility": "-"},
            {"name": "brand", "type": "String", "visibility": "#"},
            {"name": "model", "type": "String", "visibility": "#"},
            {"name": "year", "type": "int", "visibility": "#"},
            {"name": "price", "type": "double", "visibility": "+"},
            {"name": "vehicleCount", "type": "int", "visibility": "+", "static": true}
          ],
          "methods": [
            {"name": "Vehicle", "parameters": [{"name": "brand", "type": "String"}, {"name": "model", "type": "String"}], "return_type": "Vehicle", "visibility": "+"},
            {"name": "startEngine", "return_type": "void", "visibility": "+", "abstract": true},
            {"name": "stopEngine", "return_type": "void", "visibility": "+"},
            {"name": "getPrice", "return_type": "double", "visibility": "+"},
            {"name": "setPrice", "parameters": [{"name": "price", "type": "double"}], "return_type": "void", "visibility": "+"},
            {"name": "getTotalVehicles", "return_type": "int", "visibility": "+", "static": true}
          ]
        },
        {
          "name": "Car",
          "type": "class",
          "attributes": [
            {"name": "doors", "type": "int", "visibility": "-"},
            {"name": "transmission", "type": "TransmissionType", "visibility": "-"},
            {"name": "fuelType", "type": "FuelType", "visibility": "-"}
          ],
          "methods": [
            {"name": "Car", "parameters": [{"name": "brand", "type": "String"}, {"name": "model", "type": "String"}, {"name": "doors", "type": "int"}], "return_type": "Car", "visibility": "+"},
            {"name": "startEngine", "return_type": "void", "visibility": "+"},
            {"name": "openTrunk", "return_type": "void", "visibility": "+"},
            {"name": "getDoors", "return_type": "int", "visibility": "+"},
            {"name": "setTransmission", "parameters": [{"name": "transmission", "type": "TransmissionType"}], "return_type": "void", "visibility": "+"}
          ]
        },
        {
          "name": "Motorcycle",
          "type": "class",
          "attributes": [
            {"name": "engineSize", "type": "int", "visibility": "-"},
            {"name": "hasWindshield", "type": "boolean", "visibility": "-"}
          ],
          "methods": [
            {"name": "Motorcycle", "parameters": [{"name": "brand", "type": "String"}, {"name": "model", "type": "String"}], "return_type": "Motorcycle", "visibility": "+"},
            {"name": "startEngine", "return_type": "void", "visibility": "+"},
            {"name": "wheelie", "return_type": "void", "visibility": "+"},
            {"name": "getEngineSize", "return_type": "int", "visibility": "+"}
          ]
        },
        {
          "name": "Engine",
          "type": "class",
          "attributes": [
            {"name": "horsepower", "type": "int", "visibility": "-"},
            {"name": "cylinders", "type": "int", "visibility": "-"},
            {"name": "fuelType", "type": "FuelType", "visibility": "-"}
          ],
          "methods": [
            {"name": "Engine", "parameters": [{"name": "horsepower", "type": "int"}, {"name": "cylinders", "type": "int"}], "return_type": "Engine", "visibility": "+"},
            {"name": "start", "return_type": "boolean", "visibility": "+"},
            {"name": "stop", "return_type": "void", "visibility": "+"},
            {"name": "getHorsepower", "return_type": "int", "visibility": "+"}
          ]
        },
        {
          "name": "TransmissionType",
          "type": "enum",
          "attributes": [
            {"name": "MANUAL", "type": "TransmissionType", "visibility": "+", "static": true},
            {"name": "AUTOMATIC", "type": "TransmissionType", "visibility": "+", "static": true},
            {"name": "CVT", "type": "TransmissionType", "visibility": "+", "static": true}
          ],
          "methods": []
        },
        {
          "name": "FuelType",
          "type": "enum",
          "attributes": [
            {"name": "GASOLINE", "type": "FuelType", "visibility": "+", "static": true},
            {"name": "DIESEL", "type": "FuelType", "visibility": "+", "static": true},
            {"name": "ELECTRIC", "type": "FuelType", "visibility": "+", "static": true},
            {"name": "HYBRID", "type": "FuelType", "visibility": "+", "static": true}
          ],
          "methods": []
        },
        {
          "name": "VehicleService",
          "type": "interface",
          "attributes": [],
          "methods": [
            {"name": "maintenance", "parameters": [{"name": "vehicle", "type": "Vehicle"}], "return_type": "void", "visibility": "+", "abstract": true},
            {"name": "repair", "parameters": [{"name": "vehicle", "type": "Vehicle"}, {"name": "issue", "type": "String"}], "return_type": "boolean", "visibility": "+", "abstract": true},
            {"name": "inspectVehicle", "parameters": [{"name": "vehicle", "type": "Vehicle"}], "return_type": "InspectionReport", "visibility": "+", "abstract": true}
          ]
        },
        {
          "name": "GarageService",
          "type": "class",
          "attributes": [
            {"name": "garageName", "type": "String", "visibility": "-"},
            {"name": "location", "type": "String", "visibility": "-"}
          ],
          "methods": [
            {"name": "GarageService", "parameters": [{"name": "name", "type": "String"}], "return_type": "GarageService", "visibility": "+"},
            {"name": "maintenance", "parameters": [{"name": "vehicle", "type": "Vehicle"}], "return_type": "void", "visibility": "+"},
            {"name": "repair", "parameters": [{"name": "vehicle", "type": "Vehicle"}, {"name": "issue", "type": "String"}], "return_type": "boolean", "visibility": "+"},
            {"name": "inspectVehicle", "parameters": [{"name": "vehicle", "type": "Vehicle"}], "return_type": "InspectionReport", "visibility": "+"}
          ]
        }
      ],
      "relationships": [
        {
          "from": "Car",
          "to": "Vehicle",
          "type": "inheritance",
          "label": "extends"
        },
        {
          "from": "Motorcycle",
          "to": "Vehicle",
          "type": "inheritance",
          "label": "extends"
        },
        {
          "from": "Car",
          "to": "Engine",
          "type": "composition",
          "label": "has",
          "multiplicity_from": "1",
          "multiplicity_to": "1"
        },
        {
          "from": "Motorcycle",
          "to": "Engine",
          "type": "composition",
          "label": "has",
          "multiplicity_from": "1",
          "multiplicity_to": "1"
        },
        {
          "from": "Car",
          "to": "TransmissionType",
          "type": "association",
          "label": "uses",
          "multiplicity_from": "1",
          "multiplicity_to": "1"
        },
        {
          "from": "Vehicle",
          "to": "FuelType",
          "type": "association",
          "label": "uses",
          "multiplicity_from": "1",
          "multiplicity_to": "1"
        },
        {
          "from": "GarageService",
          "to": "VehicleService",
          "type": "realization",
          "label": "implements"
        },
        {
          "from": "GarageService",
          "to": "Vehicle",
          "type": "dependency",
          "label": "services",
          "multiplicity_from": "1",
          "multiplicity_to": "*"
        }
      ]
    }

    Returns:
        str: The filepath to the generated PNG image file.
    """
    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.0',
                'ranksep': '1.5'
            }
        )
        
        classes = data.get('classes', [])
        relationships = data.get('relationships', [])
        
        for cls in 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}")
            
            class_parts = []
            
            header = class_name
            if class_type == 'abstract':
                header = f"<<abstract>>\\n{class_name}"
            elif class_type == 'interface':
                header = f"<<interface>>\\n{class_name}"
            elif class_type == 'enum':
                header = f"<<enumeration>>\\n{class_name}"
            
            class_parts.append(header)
            
            if attributes:
                attr_section = ""
                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}"
                    if attr_type:
                        attr_line += f" : {attr_type}"
                    if is_static:
                        attr_line += " [static]"
                    
                    if attr_section:
                        attr_section += "\\l"
                    attr_section += attr_line
                
                if attr_section:
                    attr_section += "\\l"
                    class_parts.append(attr_section)
            
            if methods:
                method_section = ""
                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} {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}"
                    
                    if is_static:
                        method_line += " [static]"
                    if is_abstract:
                        method_line += " [abstract]"
                    
                    if method_section:
                        method_section += "\\l"
                    method_section += method_line
                
                if method_section:
                    method_section += "\\l"
                    class_parts.append(method_section)
            
            class_label = "|".join(class_parts)
            
            if class_type == 'interface':
                style = 'filled,dashed'
                fillcolor = '#f5f5f5'
            elif class_type == 'abstract':
                style = 'filled'
                fillcolor = '#eeeeee'
            else:
                style = 'filled'
                fillcolor = 'white'
            
            dot.node(
                class_name,
                class_label,
                shape='record',
                style=style,
                fillcolor=fillcolor,
                color='black'
            )
        
        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_attrs = {'color': 'black'}
            
            if label:
                edge_attrs['label'] = label
            
            if multiplicity_from:
                edge_attrs['taillabel'] = multiplicity_from
            
            if multiplicity_to:
                edge_attrs['headlabel'] = multiplicity_to
            
            if rel_type == 'inheritance':
                edge_attrs['arrowhead'] = 'empty'
            elif rel_type == 'composition':
                edge_attrs['arrowhead'] = 'normal'
                edge_attrs['arrowtail'] = 'diamond'
                edge_attrs['dir'] = 'both'
            elif rel_type == 'aggregation':
                edge_attrs['arrowhead'] = 'normal'
                edge_attrs['arrowtail'] = 'odiamond'
                edge_attrs['dir'] = 'both'
            elif rel_type == 'realization':
                edge_attrs['arrowhead'] = 'empty'
                edge_attrs['style'] = 'dashed'
            elif rel_type == 'dependency':
                edge_attrs['arrowhead'] = 'normal'
                edge_attrs['style'] = 'dashed'
            else:
                edge_attrs['arrowhead'] = 'normal'
            
            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)
            return f"{tmp.name}.{output_format}"

    except json.JSONDecodeError:
        return "Error: Invalid JSON format"
    except Exception as e:
        return f"Error: {str(e)}"