|
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)}" |