dot_printer.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. # Copyright (c) 2021 Pierre Sassoulas <pierre.sassoulas@gmail.com>
  2. # Copyright (c) 2021 Daniël van Noord <13665637+DanielNoord@users.noreply.github.com>
  3. # Copyright (c) 2021 Ashley Whetter <ashley@awhetter.co.uk>
  4. # Copyright (c) 2021 Nick Drozd <nicholasdrozd@gmail.com>
  5. # Copyright (c) 2021 Marc Mueller <30130371+cdce8p@users.noreply.github.com>
  6. # Copyright (c) 2021 Andreas Finkler <andi.finkler@gmail.com>
  7. # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
  8. # For details: https://github.com/PyCQA/pylint/blob/main/LICENSE
  9. """
  10. Class to generate files in dot format and image formats supported by Graphviz.
  11. """
  12. import os
  13. import subprocess
  14. import sys
  15. import tempfile
  16. from pathlib import Path
  17. from typing import Dict, FrozenSet, List, Optional
  18. from astroid import nodes
  19. from pylint.pyreverse.printer import EdgeType, Layout, NodeProperties, NodeType, Printer
  20. from pylint.pyreverse.utils import check_graphviz_availability, get_annotation_label
  21. ALLOWED_CHARSETS: FrozenSet[str] = frozenset(("utf-8", "iso-8859-1", "latin1"))
  22. SHAPES: Dict[NodeType, str] = {
  23. NodeType.PACKAGE: "box",
  24. NodeType.INTERFACE: "record",
  25. NodeType.CLASS: "record",
  26. }
  27. ARROWS: Dict[EdgeType, Dict[str, str]] = {
  28. EdgeType.INHERITS: dict(arrowtail="none", arrowhead="empty"),
  29. EdgeType.IMPLEMENTS: dict(arrowtail="node", arrowhead="empty", style="dashed"),
  30. EdgeType.ASSOCIATION: dict(
  31. fontcolor="green", arrowtail="none", arrowhead="diamond", style="solid"
  32. ),
  33. EdgeType.USES: dict(arrowtail="none", arrowhead="open"),
  34. }
  35. class DotPrinter(Printer):
  36. DEFAULT_COLOR = "black"
  37. def __init__(
  38. self,
  39. title: str,
  40. layout: Optional[Layout] = None,
  41. use_automatic_namespace: Optional[bool] = None,
  42. ):
  43. layout = layout or Layout.BOTTOM_TO_TOP
  44. self.charset = "utf-8"
  45. super().__init__(title, layout, use_automatic_namespace)
  46. def _open_graph(self) -> None:
  47. """Emit the header lines"""
  48. self.emit(f'digraph "{self.title}" {{')
  49. if self.layout:
  50. self.emit(f"rankdir={self.layout.value}")
  51. if self.charset:
  52. assert (
  53. self.charset.lower() in ALLOWED_CHARSETS
  54. ), f"unsupported charset {self.charset}"
  55. self.emit(f'charset="{self.charset}"')
  56. def emit_node(
  57. self,
  58. name: str,
  59. type_: NodeType,
  60. properties: Optional[NodeProperties] = None,
  61. ) -> None:
  62. """Create a new node. Nodes can be classes, packages, participants etc."""
  63. if properties is None:
  64. properties = NodeProperties(label=name)
  65. shape = SHAPES[type_]
  66. color = properties.color if properties.color is not None else self.DEFAULT_COLOR
  67. style = "filled" if color != self.DEFAULT_COLOR else "solid"
  68. label = self._build_label_for_node(properties)
  69. label_part = f', label="{label}"' if label else ""
  70. fontcolor_part = (
  71. f', fontcolor="{properties.fontcolor}"' if properties.fontcolor else ""
  72. )
  73. self.emit(
  74. f'"{name}" [color="{color}"{fontcolor_part}{label_part}, shape="{shape}", style="{style}"];'
  75. )
  76. def _build_label_for_node(
  77. self, properties: NodeProperties, is_interface: Optional[bool] = False
  78. ) -> str:
  79. if not properties.label:
  80. return ""
  81. label: str = properties.label
  82. if is_interface:
  83. # add a stereotype
  84. label = "<<interface>>\\n" + label
  85. if properties.attrs is None and properties.methods is None:
  86. # return a "compact" form which only displays the class name in a box
  87. return label
  88. # Add class attributes
  89. attrs: List[str] = properties.attrs or []
  90. label = "{" + label + "|" + r"\l".join(attrs) + r"\l|"
  91. # Add class methods
  92. methods: List[nodes.FunctionDef] = properties.methods or []
  93. for func in methods:
  94. args = self._get_method_arguments(func)
  95. label += fr"{func.name}({', '.join(args)})"
  96. if func.returns:
  97. label += ": " + get_annotation_label(func.returns)
  98. label += r"\l"
  99. label += "}"
  100. return label
  101. def emit_edge(
  102. self,
  103. from_node: str,
  104. to_node: str,
  105. type_: EdgeType,
  106. label: Optional[str] = None,
  107. ) -> None:
  108. """Create an edge from one node to another to display relationships."""
  109. arrowstyle = ARROWS[type_]
  110. attrs = [f'{prop}="{value}"' for prop, value in arrowstyle.items()]
  111. if label:
  112. attrs.append(f'label="{label}"')
  113. self.emit(f'"{from_node}" -> "{to_node}" [{", ".join(sorted(attrs))}];')
  114. def generate(self, outputfile: str) -> None:
  115. self._close_graph()
  116. graphviz_extensions = ("dot", "gv")
  117. name = self.title
  118. if outputfile is None:
  119. target = "png"
  120. pdot, dot_sourcepath = tempfile.mkstemp(".gv", name)
  121. ppng, outputfile = tempfile.mkstemp(".png", name)
  122. os.close(pdot)
  123. os.close(ppng)
  124. else:
  125. target = Path(outputfile).suffix.lstrip(".")
  126. if not target:
  127. target = "png"
  128. outputfile = outputfile + "." + target
  129. if target not in graphviz_extensions:
  130. pdot, dot_sourcepath = tempfile.mkstemp(".gv", name)
  131. os.close(pdot)
  132. else:
  133. dot_sourcepath = outputfile
  134. with open(dot_sourcepath, "w", encoding="utf8") as outfile:
  135. outfile.writelines(self.lines)
  136. if target not in graphviz_extensions:
  137. check_graphviz_availability()
  138. use_shell = sys.platform == "win32"
  139. subprocess.call(
  140. ["dot", "-T", target, dot_sourcepath, "-o", outputfile],
  141. shell=use_shell,
  142. )
  143. os.unlink(dot_sourcepath)
  144. def _close_graph(self) -> None:
  145. """Emit the lines needed to properly close the graph."""
  146. self.emit("}\n")