ZahirJS commited on
Commit
dfe2cff
·
verified ·
1 Parent(s): 5230056

Update class_diagram_generator.py

Browse files
Files changed (1) hide show
  1. class_diagram_generator.py +37 -43
class_diagram_generator.py CHANGED
@@ -208,17 +208,7 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
208
  'bgcolor': 'white',
209
  'pad': '0.5',
210
  'nodesep': '1.0',
211
- 'ranksep': '1.5',
212
- 'fontname': 'Arial',
213
- 'dpi': '300'
214
- },
215
- node_attr={
216
- 'fontname': 'Arial',
217
- 'fontsize': '10'
218
- },
219
- edge_attr={
220
- 'fontname': 'Arial',
221
- 'fontsize': '9'
222
  }
223
  )
224
 
@@ -234,19 +224,20 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
234
  if not class_name:
235
  raise ValueError(f"Invalid class: {cls}")
236
 
237
- class_label = ""
238
 
 
239
  if class_type == 'abstract':
240
- class_label += "<<abstract>>\\n"
241
  elif class_type == 'interface':
242
- class_label += "<<interface>>\\n"
243
  elif class_type == 'enum':
244
- class_label += "<<enumeration>>\\n"
245
 
246
- class_label += f"{class_name}"
247
 
248
  if attributes:
249
- class_label += "|"
250
  for attr in attributes:
251
  visibility = attr.get('visibility', '+')
252
  name = attr.get('name', '')
@@ -257,14 +248,18 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
257
  if attr_type:
258
  attr_line += f" : {attr_type}"
259
  if is_static:
260
- attr_line = f"{attr_line} (static)"
261
- class_label += f"{attr_line}\\l"
 
 
 
 
 
 
 
262
 
263
  if methods:
264
- if not attributes:
265
- class_label += "|"
266
- else:
267
- class_label += "|"
268
  for method in methods:
269
  visibility = method.get('visibility', '+')
270
  name = method.get('name', '')
@@ -284,18 +279,26 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
284
  method_line += f") : {return_type}"
285
 
286
  if is_static:
287
- method_line = f"{method_line} (static)"
288
  if is_abstract:
289
- method_line = f"{method_line} (abstract)"
290
 
291
- class_label += f"{method_line}\\l"
 
 
 
 
 
 
 
 
292
 
293
  if class_type == 'interface':
294
  style = 'filled,dashed'
295
- fillcolor = '#f0f0f0'
296
  elif class_type == 'abstract':
297
  style = 'filled'
298
- fillcolor = '#e8e8e8'
299
  else:
300
  style = 'filled'
301
  fillcolor = 'white'
@@ -306,8 +309,7 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
306
  shape='record',
307
  style=style,
308
  fillcolor=fillcolor,
309
- color='black',
310
- penwidth='1'
311
  )
312
 
313
  for relationship in relationships:
@@ -321,43 +323,35 @@ def generate_class_diagram(json_input: str, output_format: str) -> str:
321
  if not all([from_class, to_class]):
322
  raise ValueError(f"Invalid relationship: {relationship}")
323
 
324
- edge_attrs = {
325
- 'color': 'black',
326
- 'fontcolor': 'black',
327
- 'fontsize': '9'
328
- }
329
 
330
  if label:
331
  edge_attrs['label'] = label
332
 
333
- if multiplicity_from and multiplicity_to:
334
- edge_attrs['headlabel'] = multiplicity_to
335
  edge_attrs['taillabel'] = multiplicity_from
336
 
 
 
 
337
  if rel_type == 'inheritance':
338
  edge_attrs['arrowhead'] = 'empty'
339
- edge_attrs['arrowsize'] = '1.2'
340
  elif rel_type == 'composition':
341
  edge_attrs['arrowhead'] = 'normal'
342
  edge_attrs['arrowtail'] = 'diamond'
343
  edge_attrs['dir'] = 'both'
344
- edge_attrs['arrowsize'] = '1.0'
345
  elif rel_type == 'aggregation':
346
  edge_attrs['arrowhead'] = 'normal'
347
  edge_attrs['arrowtail'] = 'odiamond'
348
  edge_attrs['dir'] = 'both'
349
- edge_attrs['arrowsize'] = '1.0'
350
  elif rel_type == 'realization':
351
  edge_attrs['arrowhead'] = 'empty'
352
  edge_attrs['style'] = 'dashed'
353
- edge_attrs['arrowsize'] = '1.2'
354
  elif rel_type == 'dependency':
355
  edge_attrs['arrowhead'] = 'normal'
356
  edge_attrs['style'] = 'dashed'
357
- edge_attrs['arrowsize'] = '1.0'
358
  else:
359
  edge_attrs['arrowhead'] = 'normal'
360
- edge_attrs['arrowsize'] = '1.0'
361
 
362
  dot.edge(from_class, to_class, **edge_attrs)
363
 
 
208
  'bgcolor': 'white',
209
  'pad': '0.5',
210
  'nodesep': '1.0',
211
+ 'ranksep': '1.5'
 
 
 
 
 
 
 
 
 
 
212
  }
213
  )
214
 
 
224
  if not class_name:
225
  raise ValueError(f"Invalid class: {cls}")
226
 
227
+ class_parts = []
228
 
229
+ header = class_name
230
  if class_type == 'abstract':
231
+ header = f"<<abstract>>\\n{class_name}"
232
  elif class_type == 'interface':
233
+ header = f"<<interface>>\\n{class_name}"
234
  elif class_type == 'enum':
235
+ header = f"<<enumeration>>\\n{class_name}"
236
 
237
+ class_parts.append(header)
238
 
239
  if attributes:
240
+ attr_section = ""
241
  for attr in attributes:
242
  visibility = attr.get('visibility', '+')
243
  name = attr.get('name', '')
 
248
  if attr_type:
249
  attr_line += f" : {attr_type}"
250
  if is_static:
251
+ attr_line += " [static]"
252
+
253
+ if attr_section:
254
+ attr_section += "\\l"
255
+ attr_section += attr_line
256
+
257
+ if attr_section:
258
+ attr_section += "\\l"
259
+ class_parts.append(attr_section)
260
 
261
  if methods:
262
+ method_section = ""
 
 
 
263
  for method in methods:
264
  visibility = method.get('visibility', '+')
265
  name = method.get('name', '')
 
279
  method_line += f") : {return_type}"
280
 
281
  if is_static:
282
+ method_line += " [static]"
283
  if is_abstract:
284
+ method_line += " [abstract]"
285
 
286
+ if method_section:
287
+ method_section += "\\l"
288
+ method_section += method_line
289
+
290
+ if method_section:
291
+ method_section += "\\l"
292
+ class_parts.append(method_section)
293
+
294
+ class_label = "|".join(class_parts)
295
 
296
  if class_type == 'interface':
297
  style = 'filled,dashed'
298
+ fillcolor = '#f5f5f5'
299
  elif class_type == 'abstract':
300
  style = 'filled'
301
+ fillcolor = '#eeeeee'
302
  else:
303
  style = 'filled'
304
  fillcolor = 'white'
 
309
  shape='record',
310
  style=style,
311
  fillcolor=fillcolor,
312
+ color='black'
 
313
  )
314
 
315
  for relationship in relationships:
 
323
  if not all([from_class, to_class]):
324
  raise ValueError(f"Invalid relationship: {relationship}")
325
 
326
+ edge_attrs = {'color': 'black'}
 
 
 
 
327
 
328
  if label:
329
  edge_attrs['label'] = label
330
 
331
+ if multiplicity_from:
 
332
  edge_attrs['taillabel'] = multiplicity_from
333
 
334
+ if multiplicity_to:
335
+ edge_attrs['headlabel'] = multiplicity_to
336
+
337
  if rel_type == 'inheritance':
338
  edge_attrs['arrowhead'] = 'empty'
 
339
  elif rel_type == 'composition':
340
  edge_attrs['arrowhead'] = 'normal'
341
  edge_attrs['arrowtail'] = 'diamond'
342
  edge_attrs['dir'] = 'both'
 
343
  elif rel_type == 'aggregation':
344
  edge_attrs['arrowhead'] = 'normal'
345
  edge_attrs['arrowtail'] = 'odiamond'
346
  edge_attrs['dir'] = 'both'
 
347
  elif rel_type == 'realization':
348
  edge_attrs['arrowhead'] = 'empty'
349
  edge_attrs['style'] = 'dashed'
 
350
  elif rel_type == 'dependency':
351
  edge_attrs['arrowhead'] = 'normal'
352
  edge_attrs['style'] = 'dashed'
 
353
  else:
354
  edge_attrs['arrowhead'] = 'normal'
 
355
 
356
  dot.edge(from_class, to_class, **edge_attrs)
357