graph.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212
  1. # Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
  2. # For details: https://github.com/pylint-dev/pylint/blob/main/LICENSE
  3. # Copyright (c) https://github.com/pylint-dev/pylint/blob/main/CONTRIBUTORS.txt
  4. """Graph manipulation utilities.
  5. (dot generation adapted from pypy/translator/tool/make_dot.py)
  6. """
  7. from __future__ import annotations
  8. import codecs
  9. import os
  10. import shutil
  11. import subprocess
  12. import tempfile
  13. from collections.abc import Sequence
  14. from typing import Any
  15. def target_info_from_filename(filename: str) -> tuple[str, str, str]:
  16. """Transforms /some/path/foo.png into ('/some/path', 'foo.png', 'png')."""
  17. basename = os.path.basename(filename)
  18. storedir = os.path.dirname(os.path.abspath(filename))
  19. target = os.path.splitext(filename)[-1][1:]
  20. return storedir, basename, target
  21. class DotBackend:
  22. """Dot File back-end."""
  23. def __init__(
  24. self,
  25. graphname: str,
  26. rankdir: str | None = None,
  27. size: Any = None,
  28. ratio: Any = None,
  29. charset: str = "utf-8",
  30. renderer: str = "dot",
  31. additional_param: dict[str, Any] | None = None,
  32. ) -> None:
  33. if additional_param is None:
  34. additional_param = {}
  35. self.graphname = graphname
  36. self.renderer = renderer
  37. self.lines: list[str] = []
  38. self._source: str | None = None
  39. self.emit(f"digraph {normalize_node_id(graphname)} {{")
  40. if rankdir:
  41. self.emit(f"rankdir={rankdir}")
  42. if ratio:
  43. self.emit(f"ratio={ratio}")
  44. if size:
  45. self.emit(f'size="{size}"')
  46. if charset:
  47. assert charset.lower() in {
  48. "utf-8",
  49. "iso-8859-1",
  50. "latin1",
  51. }, f"unsupported charset {charset}"
  52. self.emit(f'charset="{charset}"')
  53. for param in additional_param.items():
  54. self.emit("=".join(param))
  55. def get_source(self) -> str:
  56. """Returns self._source."""
  57. if self._source is None:
  58. self.emit("}\n")
  59. self._source = "\n".join(self.lines)
  60. del self.lines
  61. return self._source
  62. source = property(get_source)
  63. def generate(
  64. self, outputfile: str | None = None, mapfile: str | None = None
  65. ) -> str:
  66. """Generates a graph file.
  67. :param str outputfile: filename and path [defaults to graphname.png]
  68. :param str mapfile: filename and path
  69. :rtype: str
  70. :return: a path to the generated file
  71. :raises RuntimeError: if the executable for rendering was not found
  72. """
  73. # pylint: disable=duplicate-code
  74. graphviz_extensions = ("dot", "gv")
  75. name = self.graphname
  76. if outputfile is None:
  77. target = "png"
  78. pdot, dot_sourcepath = tempfile.mkstemp(".gv", name)
  79. ppng, outputfile = tempfile.mkstemp(".png", name)
  80. os.close(pdot)
  81. os.close(ppng)
  82. else:
  83. _, _, target = target_info_from_filename(outputfile)
  84. if not target:
  85. target = "png"
  86. outputfile = outputfile + "." + target
  87. if target not in graphviz_extensions:
  88. pdot, dot_sourcepath = tempfile.mkstemp(".gv", name)
  89. os.close(pdot)
  90. else:
  91. dot_sourcepath = outputfile
  92. with codecs.open(dot_sourcepath, "w", encoding="utf8") as file:
  93. file.write(self.source)
  94. if target not in graphviz_extensions:
  95. if shutil.which(self.renderer) is None:
  96. raise RuntimeError(
  97. f"Cannot generate `{outputfile}` because '{self.renderer}' "
  98. "executable not found. Install graphviz, or specify a `.gv` "
  99. "outputfile to produce the DOT source code."
  100. )
  101. if mapfile:
  102. subprocess.run(
  103. [
  104. self.renderer,
  105. "-Tcmapx",
  106. "-o",
  107. mapfile,
  108. "-T",
  109. target,
  110. dot_sourcepath,
  111. "-o",
  112. outputfile,
  113. ],
  114. check=True,
  115. )
  116. else:
  117. subprocess.run(
  118. [self.renderer, "-T", target, dot_sourcepath, "-o", outputfile],
  119. check=True,
  120. )
  121. os.unlink(dot_sourcepath)
  122. return outputfile
  123. def emit(self, line: str) -> None:
  124. """Adds <line> to final output."""
  125. self.lines.append(line)
  126. def emit_edge(self, name1: str, name2: str, **props: Any) -> None:
  127. """Emit an edge from <name1> to <name2>.
  128. For edge properties: see https://www.graphviz.org/doc/info/attrs.html
  129. """
  130. attrs = [f'{prop}="{value}"' for prop, value in props.items()]
  131. n_from, n_to = normalize_node_id(name1), normalize_node_id(name2)
  132. self.emit(f"{n_from} -> {n_to} [{', '.join(sorted(attrs))}];")
  133. def emit_node(self, name: str, **props: Any) -> None:
  134. """Emit a node with given properties.
  135. For node properties: see https://www.graphviz.org/doc/info/attrs.html
  136. """
  137. attrs = [f'{prop}="{value}"' for prop, value in props.items()]
  138. self.emit(f"{normalize_node_id(name)} [{', '.join(sorted(attrs))}];")
  139. def normalize_node_id(nid: str) -> str:
  140. """Returns a suitable DOT node id for `nid`."""
  141. return f'"{nid}"'
  142. def get_cycles(
  143. graph_dict: dict[str, set[str]], vertices: list[str] | None = None
  144. ) -> Sequence[list[str]]:
  145. """Return a list of detected cycles based on an ordered graph (i.e. keys are
  146. vertices and values are lists of destination vertices representing edges).
  147. """
  148. if not graph_dict:
  149. return ()
  150. result: list[list[str]] = []
  151. if vertices is None:
  152. vertices = list(graph_dict.keys())
  153. for vertice in vertices:
  154. _get_cycles(graph_dict, [], set(), result, vertice)
  155. return result
  156. def _get_cycles(
  157. graph_dict: dict[str, set[str]],
  158. path: list[str],
  159. visited: set[str],
  160. result: list[list[str]],
  161. vertice: str,
  162. ) -> None:
  163. """Recursive function doing the real work for get_cycles."""
  164. if vertice in path:
  165. cycle = [vertice]
  166. for node in path[::-1]:
  167. if node == vertice:
  168. break
  169. cycle.insert(0, node)
  170. # make a canonical representation
  171. start_from = min(cycle)
  172. index = cycle.index(start_from)
  173. cycle = cycle[index:] + cycle[0:index]
  174. # append it to result if not already in
  175. if cycle not in result:
  176. result.append(cycle)
  177. return
  178. path.append(vertice)
  179. try:
  180. for node in graph_dict[vertice]:
  181. # don't check already visited nodes again
  182. if node not in visited:
  183. _get_cycles(graph_dict, path, visited, result, node)
  184. visited.add(node)
  185. except KeyError:
  186. pass
  187. path.pop()