Create common Printer base class for pyreverse and improve typing.

pylint/pyreverse/diagrams.py

@@ -45,6 +45,19 @@ def __init__(self, title="No name", node=None):
         self.node = node
 
 
+class PackageEntity(DiagramEntity):
+    """A diagram object representing a package"""
+
+
+class ClassEntity(DiagramEntity):
+    """A diagram object representing a class"""
+
+    def __init__(self, title, node):
+        super().__init__(title=title, node=node)
+        self.attrs = None
+        self.methods = None
+
+
 class ClassDiagram(Figure, FilterMixIn):
     """main class diagram handling"""
 

pylint/pyreverse/dot_printer.py

@@ -0,0 +1,141 @@
+# Copyright (c) 2021 Andreas Finkler <[email protected]>
+
+# Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
+# For details: https://github.com/PyCQA/pylint/blob/main/LICENSE
+
+"""
+Class to generate files in dot format and image formats supported by Graphviz.
+"""
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Dict, FrozenSet, Optional
+
+from pylint.pyreverse.printer import EdgeType, Layout, NodeProperties, NodeType, Printer
+
+ALLOWED_CHARSETS: FrozenSet[str] = frozenset(("utf-8", "iso-8859-1", "latin1"))
+SHAPES: Dict[NodeType, str] = {
+    NodeType.PACKAGE: "box",
+    NodeType.INTERFACE: "record",
+    NodeType.CLASS: "record",
+}
+ARROWS: Dict[EdgeType, Dict] = {
+    EdgeType.INHERITS: dict(arrowtail="none", arrowhead="empty"),
+    EdgeType.IMPLEMENTS: dict(arrowtail="node", arrowhead="empty", style="dashed"),
+    EdgeType.ASSOCIATION: dict(
+        fontcolor="green", arrowtail="none", arrowhead="diamond", style="solid"
+    ),
+    EdgeType.USES: dict(arrowtail="none", arrowhead="open"),
+}
+RANKDIR: Dict[Layout, str] = {
+    Layout.LEFT_TO_RIGHT: "LR",
+    Layout.RIGHT_TO_LEFT: "RL",
+    Layout.TOP_TO_BOTTOM: "TB",
+    Layout.BOTTOM_TO_TOP: "BT",
+}
+
+
+class DotPrinter(Printer):
+    def __init__(
+        self,
+        title: str,
+        layout: Optional[Layout] = None,
+        use_automatic_namespace: Optional[bool] = None,
+    ):
+        self.charset = "utf-8"
+        self.node_style = "solid"
+        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={RANKDIR[self.layout]}")
+        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: Optional[NodeProperties] = 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 "black"
+        label = properties.label
+        if label:
+            if type_ is NodeType.INTERFACE:
+                label = "<<interface>>\\n" + label
+            label_part = f', label="{label}"'
+        else:
+            label_part = ""
+        fontcolor_part = (
+            f', fontcolor="{properties.fontcolor}"' if properties.fontcolor else ""
+        )
+        self.emit(
+            f'"{name}" [color="{color}"{fontcolor_part}{label_part}, shape="{shape}", style="{self.node_style}"];'
+        )
+
+    def emit_edge(
+        self,
+        from_node: str,
+        to_node: str,
+        type_: EdgeType,
+        label: Optional[str] = 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:
+            if shutil.which("dot") is None:
+                raise RuntimeError(
+                    f"Cannot generate `{outputfile}` because 'dot' "
+                    "executable not found. Install graphviz, or specify a `.gv` "
+                    "outputfile to produce the DOT source code."
+                )
+            use_shell = sys.platform == "win32"
+            subprocess.call(
+                ["dot", "-T", target, dot_sourcepath, "-o", outputfile],
+                shell=use_shell,
+            )
+            os.unlink(dot_sourcepath)
+        # return outputfile  TODO should we return this?
+
+    def _close_graph(self) -> None:
+        """Emit the lines needed to properly close the graph."""
+        self.emit("}\n")

pylint/pyreverse/printer.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2021 Andreas Finkler <[email protected]>
+
+# Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
+# For details: https://github.com/PyCQA/pylint/blob/main/LICENSE
+
+"""
+Base class defining the interface for a printer.
+"""
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import List, NamedTuple, Optional
+
+
+class NodeType(Enum):
+    CLASS = "class"
+    INTERFACE = "interface"
+    PACKAGE = "package"
+
+
+class EdgeType(Enum):
+    INHERITS = "inherits"
+    IMPLEMENTS = "implements"
+    ASSOCIATION = "association"
+    USES = "uses"
+
+
+class Layout(Enum):
+    LEFT_TO_RIGHT = "LR"
+    RIGHT_TO_LEFT = "RL"
+    TOP_TO_BOTTOM = "TB"
+    BOTTOM_TO_TOP = "BT"
+
+
+class NodeProperties(NamedTuple):
+    label: str
+    color: Optional[str] = None
+    fontcolor: Optional[str] = None
+    body: Optional[str] = None
+
+
+class Printer(ABC):
+    """Base class defining the interface for a printer"""
+
+    def __init__(
+        self,
+        title: str,
+        layout: Optional[Layout] = None,
+        use_automatic_namespace: Optional[bool] = None,
+    ):
+        self.title: str = title
+        self.layout = layout
+        self.use_automatic_namespace = use_automatic_namespace
+        self.lines: List[str] = []
+        self._open_graph()
+
+    @abstractmethod
+    def _open_graph(self) -> None:
+        """Emit the header lines, i.e. all boilerplate code that defines things like layout etc."""
+
+    def emit(self, line: str, force_newline: Optional[bool] = True) -> None:
+        if force_newline and not line.endswith("\n"):
+            line += "\n"
+        self.lines.append(line)
+
+    @abstractmethod
+    def emit_node(
+        self,
+        name: str,
+        type_: NodeType,
+        properties: Optional[NodeProperties] = None,
+    ) -> None:
+        """Create a new node. Nodes can be classes, packages, participants etc."""
+
+    @abstractmethod
+    def emit_edge(
+        self,
+        from_node: str,
+        to_node: str,
+        type_: EdgeType,
+        label: Optional[str] = None,
+    ) -> None:
+        """Create an edge from one node to another to display relationships."""
+
+    def generate(self, outputfile: str) -> None:
+        """Generate and save the final outputfile."""
+        self._close_graph()
+        with open(outputfile, "w", encoding="utf-8") as outfile:
+            outfile.writelines(self.lines)
+
+    @abstractmethod
+    def _close_graph(self) -> None:
+        """Emit the lines needed to properly close the graph."""