plantuml_printer.py 3.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. # Copyright (c) 2021 Andreas Finkler <andi.finkler@gmail.com>
  2. # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
  3. # For details: https://github.com/PyCQA/pylint/blob/main/LICENSE
  4. """
  5. Class to generate files in dot format and image formats supported by Graphviz.
  6. """
  7. from typing import Dict, Optional
  8. from pylint.pyreverse.printer import EdgeType, Layout, NodeProperties, NodeType, Printer
  9. from pylint.pyreverse.utils import get_annotation_label
  10. class PlantUmlPrinter(Printer):
  11. """Printer for PlantUML diagrams"""
  12. DEFAULT_COLOR = "black"
  13. NODES: Dict[NodeType, str] = {
  14. NodeType.CLASS: "class",
  15. NodeType.INTERFACE: "class",
  16. NodeType.PACKAGE: "package",
  17. }
  18. ARROWS: Dict[EdgeType, str] = {
  19. EdgeType.INHERITS: "--|>",
  20. EdgeType.IMPLEMENTS: "..|>",
  21. EdgeType.ASSOCIATION: "--*",
  22. EdgeType.USES: "-->",
  23. }
  24. def _open_graph(self) -> None:
  25. """Emit the header lines"""
  26. self.emit("@startuml " + self.title)
  27. if not self.use_automatic_namespace:
  28. self.emit("set namespaceSeparator none")
  29. if self.layout:
  30. if self.layout is Layout.LEFT_TO_RIGHT:
  31. self.emit("left to right direction")
  32. elif self.layout is Layout.TOP_TO_BOTTOM:
  33. self.emit("top to bottom direction")
  34. else:
  35. raise ValueError(
  36. f"Unsupported layout {self.layout}. PlantUmlPrinter only supports left to right and top to bottom layout."
  37. )
  38. def emit_node(
  39. self,
  40. name: str,
  41. type_: NodeType,
  42. properties: Optional[NodeProperties] = None,
  43. ) -> None:
  44. """Create a new node. Nodes can be classes, packages, participants etc."""
  45. if properties is None:
  46. properties = NodeProperties(label=name)
  47. stereotype = " << interface >>" if type_ is NodeType.INTERFACE else ""
  48. nodetype = self.NODES[type_]
  49. if properties.color and properties.color != self.DEFAULT_COLOR:
  50. color = f" #{properties.color}"
  51. else:
  52. color = ""
  53. body = []
  54. if properties.attrs:
  55. body.extend(properties.attrs)
  56. if properties.methods:
  57. for func in properties.methods:
  58. args = self._get_method_arguments(func)
  59. line = f"{func.name}({', '.join(args)})"
  60. if func.returns:
  61. line += " -> " + get_annotation_label(func.returns)
  62. body.append(line)
  63. label = properties.label if properties.label is not None else name
  64. if properties.fontcolor and properties.fontcolor != self.DEFAULT_COLOR:
  65. label = f"<color:{properties.fontcolor}>{label}</color>"
  66. self.emit(f'{nodetype} "{label}" as {name}{stereotype}{color} {{')
  67. self._inc_indent()
  68. for line in body:
  69. self.emit(line)
  70. self._dec_indent()
  71. self.emit("}")
  72. def emit_edge(
  73. self,
  74. from_node: str,
  75. to_node: str,
  76. type_: EdgeType,
  77. label: Optional[str] = None,
  78. ) -> None:
  79. """Create an edge from one node to another to display relationships."""
  80. edge = f"{from_node} {self.ARROWS[type_]} {to_node}"
  81. if label:
  82. edge += f" : {label}"
  83. self.emit(edge)
  84. def _close_graph(self) -> None:
  85. """Emit the lines needed to properly close the graph."""
  86. self.emit("@enduml")