mermaidjs_printer.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
  2. # For details: https://github.com/PyCQA/pylint/blob/main/LICENSE
  3. # Copyright (c) https://github.com/PyCQA/pylint/blob/main/CONTRIBUTORS.txt
  4. """Class to generate files in mermaidjs format."""
  5. from __future__ import annotations
  6. from pylint.pyreverse.printer import EdgeType, NodeProperties, NodeType, Printer
  7. from pylint.pyreverse.utils import get_annotation_label
  8. class MermaidJSPrinter(Printer):
  9. """Printer for MermaidJS diagrams."""
  10. DEFAULT_COLOR = "black"
  11. NODES: dict[NodeType, str] = {
  12. NodeType.CLASS: "class",
  13. NodeType.INTERFACE: "class",
  14. NodeType.PACKAGE: "class",
  15. }
  16. ARROWS: dict[EdgeType, str] = {
  17. EdgeType.INHERITS: "--|>",
  18. EdgeType.IMPLEMENTS: "..|>",
  19. EdgeType.ASSOCIATION: "--*",
  20. EdgeType.AGGREGATION: "--o",
  21. EdgeType.USES: "-->",
  22. }
  23. def _open_graph(self) -> None:
  24. """Emit the header lines."""
  25. self.emit("classDiagram")
  26. self._inc_indent()
  27. def emit_node(
  28. self,
  29. name: str,
  30. type_: NodeType,
  31. properties: NodeProperties | None = None,
  32. ) -> None:
  33. """Create a new node.
  34. Nodes can be classes, packages, participants etc.
  35. """
  36. # pylint: disable=duplicate-code
  37. if properties is None:
  38. properties = NodeProperties(label=name)
  39. stereotype = "~~Interface~~" if type_ is NodeType.INTERFACE else ""
  40. nodetype = self.NODES[type_]
  41. body = []
  42. if properties.attrs:
  43. body.extend(properties.attrs)
  44. if properties.methods:
  45. for func in properties.methods:
  46. args = self._get_method_arguments(func)
  47. line = f"{func.name}({', '.join(args)})"
  48. line += "*" if func.is_abstract() else ""
  49. if func.returns:
  50. line += f" {get_annotation_label(func.returns)}"
  51. body.append(line)
  52. name = name.split(".")[-1]
  53. self.emit(f"{nodetype} {name}{stereotype} {{")
  54. self._inc_indent()
  55. for line in body:
  56. self.emit(line)
  57. self._dec_indent()
  58. self.emit("}")
  59. def emit_edge(
  60. self,
  61. from_node: str,
  62. to_node: str,
  63. type_: EdgeType,
  64. label: str | None = None,
  65. ) -> None:
  66. """Create an edge from one node to another to display relationships."""
  67. from_node = from_node.split(".")[-1]
  68. to_node = to_node.split(".")[-1]
  69. edge = f"{from_node} {self.ARROWS[type_]} {to_node}"
  70. if label:
  71. edge += f" : {label}"
  72. self.emit(edge)
  73. def _close_graph(self) -> None:
  74. """Emit the lines needed to properly close the graph."""
  75. self._dec_indent()
  76. class HTMLMermaidJSPrinter(MermaidJSPrinter):
  77. """Printer for MermaidJS diagrams wrapped in a html boilerplate."""
  78. HTML_OPEN_BOILERPLATE = """<html>
  79. <body>
  80. <script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
  81. <div class="mermaid">
  82. """
  83. HTML_CLOSE_BOILERPLATE = """
  84. </div>
  85. </body>
  86. </html>
  87. """
  88. GRAPH_INDENT_LEVEL = 4
  89. def _open_graph(self) -> None:
  90. self.emit(self.HTML_OPEN_BOILERPLATE)
  91. for _ in range(self.GRAPH_INDENT_LEVEL):
  92. self._inc_indent()
  93. super()._open_graph()
  94. def _close_graph(self) -> None:
  95. for _ in range(self.GRAPH_INDENT_LEVEL):
  96. self._dec_indent()
  97. self.emit(self.HTML_CLOSE_BOILERPLATE)