1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
|
# Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
# For details: https://github.com/pylint-dev/pylint/blob/main/LICENSE
# Copyright (c) https://github.com/pylint-dev/pylint/blob/main/CONTRIBUTORS.txt
"""Class to generate files in dot format and image formats supported by Graphviz."""
from __future__ import annotations
import os
import subprocess
import tempfile
from enum import Enum
from pathlib import Path
from astroid import nodes
from pylint.pyreverse.printer import EdgeType, Layout, NodeProperties, NodeType, Printer
from pylint.pyreverse.utils import get_annotation_label
class HTMLLabels(Enum):
LINEBREAK_LEFT = '<br ALIGN="LEFT"/>'
ALLOWED_CHARSETS: frozenset[str] = frozenset(("utf-8", "iso-8859-1", "latin1"))
SHAPES: dict[NodeType, str] = {
NodeType.PACKAGE: "box",
NodeType.CLASS: "record",
}
# pylint: disable-next=consider-using-namedtuple-or-dataclass
ARROWS: dict[EdgeType, dict[str, str]] = {
EdgeType.INHERITS: {"arrowtail": "none", "arrowhead": "empty"},
EdgeType.ASSOCIATION: {
"fontcolor": "green",
"arrowtail": "none",
"arrowhead": "diamond",
"style": "solid",
},
EdgeType.AGGREGATION: {
"fontcolor": "green",
"arrowtail": "none",
"arrowhead": "odiamond",
"style": "solid",
},
EdgeType.USES: {"arrowtail": "none", "arrowhead": "open"},
}
class DotPrinter(Printer):
DEFAULT_COLOR = "black"
def __init__(
self,
title: str,
layout: Layout | None = None,
use_automatic_namespace: bool | None = None,
):
layout = layout or Layout.BOTTOM_TO_TOP
self.charset = "utf-8"
super().__init__(title, layout, use_automatic_namespace)
def _open_graph(self) -> None:
"""Emit the header lines."""
self.emit(f'digraph "{self.title}" {{')
if self.layout:
self.emit(f"rankdir={self.layout.value}")
if self.charset:
assert (
self.charset.lower() in ALLOWED_CHARSETS
), f"unsupported charset {self.charset}"
self.emit(f'charset="{self.charset}"')
def emit_node(
self,
name: str,
type_: NodeType,
properties: NodeProperties | None = None,
) -> None:
"""Create a new node.
Nodes can be classes, packages, participants etc.
"""
if properties is None:
properties = NodeProperties(label=name)
shape = SHAPES[type_]
color = properties.color if properties.color is not None else self.DEFAULT_COLOR
style = "filled" if color != self.DEFAULT_COLOR else "solid"
label = self._build_label_for_node(properties)
label_part = f", label=<{label}>" if label else ""
fontcolor_part = (
f', fontcolor="{properties.fontcolor}"' if properties.fontcolor else ""
)
self.emit(
f'"{name}" [color="{color}"{fontcolor_part}{label_part}, shape="{shape}", style="{style}"];'
)
def _build_label_for_node(self, properties: NodeProperties) -> str:
if not properties.label:
return ""
label: str = properties.label
if properties.attrs is None and properties.methods is None:
# return a "compact" form which only displays the class name in a box
return label
# Add class attributes
attrs: list[str] = properties.attrs or []
attrs_string = rf"{HTMLLabels.LINEBREAK_LEFT.value}".join(
attr.replace("|", r"\|") for attr in attrs
)
label = rf"{{{label}|{attrs_string}{HTMLLabels.LINEBREAK_LEFT.value}|"
# Add class methods
methods: list[nodes.FunctionDef] = properties.methods or []
for func in methods:
args = self._get_method_arguments(func)
method_name = (
f"<I>{func.name}</I>" if func.is_abstract() else f"{func.name}"
)
label += rf"{method_name}({', '.join(args)})"
if func.returns:
label += ": " + get_annotation_label(func.returns)
label += rf"{HTMLLabels.LINEBREAK_LEFT.value}"
label += "}"
return label
def emit_edge(
self,
from_node: str,
to_node: str,
type_: EdgeType,
label: str | None = None,
) -> None:
"""Create an edge from one node to another to display relationships."""
arrowstyle = ARROWS[type_]
attrs = [f'{prop}="{value}"' for prop, value in arrowstyle.items()]
if label:
attrs.append(f'label="{label}"')
self.emit(f'"{from_node}" -> "{to_node}" [{", ".join(sorted(attrs))}];')
def generate(self, outputfile: str) -> None:
self._close_graph()
graphviz_extensions = ("dot", "gv")
name = self.title
if outputfile is None:
target = "png"
pdot, dot_sourcepath = tempfile.mkstemp(".gv", name)
ppng, outputfile = tempfile.mkstemp(".png", name)
os.close(pdot)
os.close(ppng)
else:
target = Path(outputfile).suffix.lstrip(".")
if not target:
target = "png"
outputfile = outputfile + "." + target
if target not in graphviz_extensions:
pdot, dot_sourcepath = tempfile.mkstemp(".gv", name)
os.close(pdot)
else:
dot_sourcepath = outputfile
with open(dot_sourcepath, "w", encoding="utf8") as outfile:
outfile.writelines(self.lines)
if target not in graphviz_extensions:
subprocess.run(
["dot", "-T", target, dot_sourcepath, "-o", outputfile], check=True
)
os.unlink(dot_sourcepath)
def _close_graph(self) -> None:
"""Emit the lines needed to properly close the graph."""
self.emit("}\n")
|
