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,129 @@
+# 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 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
+from pylint.pyreverse.utils import check_graphviz_availability
+
+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"),
+}
+
+
+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={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: 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:
+            check_graphviz_availability()
+            use_shell = sys.platform == "win32"
+            subprocess.call(
+                ["dot", "-T", target, dot_sourcepath, "-o", outputfile],
+                shell=use_shell,
+            )
+            os.unlink(dot_sourcepath)
+
+    def _close_graph(self) -> None:
+        """Emit the lines needed to properly close the graph."""
+        self.emit("}\n")

pylint/pyreverse/main.py

@@ -21,15 +21,14 @@
   create UML diagrams for classes and modules in <packages>
 """
 import os
-import subprocess
 import sys
 from typing import Iterable
 
 from pylint.config import ConfigurationMixIn
 from pylint.pyreverse import writer
 from pylint.pyreverse.diadefslib import DiadefsHandler
 from pylint.pyreverse.inspector import Linker, project_from_files
-from pylint.pyreverse.utils import insert_default_options
+from pylint.pyreverse.utils import check_graphviz_availability, insert_default_options
 
 OPTIONS = (
     (
@@ -175,19 +174,6 @@
 )
 
 
-def _check_graphviz_available(output_format):
-    """check if we need graphviz for different output format"""
-    try:
-        subprocess.call(["dot", "-V"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    except OSError:
-        print(
-            "The output format '%s' is currently not available.\n"
-            "Please install 'Graphviz' to have other output formats "
-            "than 'dot' or 'vcg'." % output_format
-        )
-        sys.exit(32)
-
-
 class Run(ConfigurationMixIn):
     """base class providing common behaviour for pyreverse commands"""
 
@@ -198,7 +184,7 @@ def __init__(self, args: Iterable[str]):
         insert_default_options()
         args = self.load_command_line_configuration(args)
         if self.config.output_format not in ("dot", "vcg"):
-            _check_graphviz_available(self.config.output_format)
+            check_graphviz_availability()
 
         sys.exit(self.run(args))
 

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."""

pylint/pyreverse/utils.py

@@ -20,6 +20,7 @@
 """Generic classes/functions for pyreverse core/extensions. """
 import os
 import re
+import shutil
 import sys
 from typing import Optional, Union
 
@@ -275,3 +276,16 @@ def infer_node(node: Union[astroid.AssignAttr, astroid.AssignName]) -> set:
         return set(node.infer())
     except astroid.InferenceError:
         return {ann} if ann else set()
+
+
+def check_graphviz_availability():
+    """Check if the ``dot`` command is available on the machine.
+    This is needed if image output is desired and ``dot`` is used to convert
+    from *.dot or *.gv into the final output format."""
+    if shutil.which("dot") is None:
+        print(
+            "The requested output format is currently not available.\n"
+            "Please install 'Graphviz' to have other output formats "
+            "than 'dot' or 'vcg'."
+        )
+        sys.exit(32)