pyreverse: add PlantUML output (#4846)

* Extract helper method to get annotated arguments into ``Printer`` base class.

* Add ``Printer`` subclass for PlantUML output

* Add functional test for ``PlantUmlPrinter``

* Add tests for specific layout for ``PlantUmlPrinter``

* Extract test helper function to remove code duplication

* Add new test class to check type annotations

* Cleanup generated .puml files after tests finished

* Create a factory function to get the correct ``Printer`` class for a given filetype.

* Fix unittest after adding a new class to the test data.

* Add changelog and whatsnew entry

* Add "plantuml" as possible extension for PlantUML output

ChangeLog

@@ -12,6 +12,10 @@ Release date: TBA
 ..
   Put bug fixes that should not wait for a new minor version here
 
++ pyreverse: add output in PlantUML format.
+
+  Closes #4498
+
 * pylint does not crash with a traceback anymore when a file is problematic. It
   creates a template text file for opening an issue on the bug tracker instead.
   The linting can go on for other non problematic files instead of being impossible.

doc/whatsnew/2.10.rst

@@ -44,6 +44,8 @@ Extensions
 Other Changes
 =============
 
+* Pyreverse - add output in PlantUML format
+
 * pylint does not crash with a traceback anymore when a file is problematic. It
   creates a template text file for opening an issue on the bug tracker instead.
   The linting can go on for other non problematic files instead of being impossible.

pylint/pyreverse/dot_printer.py

@@ -80,9 +80,8 @@ def emit_node(
             f'"{name}" [color="{color}"{fontcolor_part}{label_part}, shape="{shape}", style="{self.node_style}"];'
         )
 
-    @staticmethod
     def _build_label_for_node(
-        properties: NodeProperties, is_interface: Optional[bool] = False
+        self, properties: NodeProperties, is_interface: Optional[bool] = False
     ) -> str:
         if not properties.label:
             return ""
@@ -103,31 +102,11 @@ def _build_label_for_node(
         # Add class methods
         methods: List[astroid.FunctionDef] = properties.methods or []
         for func in methods:
-            return_type = (
-                f": {get_annotation_label(func.returns)}" if func.returns else ""
-            )
-
-            if func.args.args:
-                arguments: List[astroid.AssignName] = [
-                    arg for arg in func.args.args if arg.name != "self"
-                ]
-            else:
-                arguments = []
-
-            annotations = dict(zip(arguments, func.args.annotations[1:]))
-            for arg in arguments:
-                annotation_label = ""
-                ann = annotations.get(arg)
-                if ann:
-                    annotation_label = get_annotation_label(ann)
-                annotations[arg] = annotation_label
-
-            args = ", ".join(
-                f"{arg.name}: {ann}" if ann else f"{arg.name}"
-                for arg, ann in annotations.items()
-            )
-
-            label += fr"{func.name}({args}){return_type}\l"
+            args = self._get_method_arguments(func)
+            label += fr"{func.name}({', '.join(args)})"
+            if func.returns:
+                label += ": " + get_annotation_label(func.returns)
+            label += r"\l"
         label += "}"
         return label
 

pylint/pyreverse/main.py

@@ -27,10 +27,8 @@
 from pylint.config import ConfigurationMixIn
 from pylint.pyreverse import writer
 from pylint.pyreverse.diadefslib import DiadefsHandler
-from pylint.pyreverse.dot_printer import DotPrinter
 from pylint.pyreverse.inspector import Linker, project_from_files
 from pylint.pyreverse.utils import check_graphviz_availability, insert_default_options
-from pylint.pyreverse.vcg_printer import VCGPrinter
 
 OPTIONS = (
     (
@@ -209,8 +207,7 @@ def run(self, args):
             diadefs = handler.get_diadefs(project, linker)
         finally:
             sys.path.pop(0)
-        printer_class = VCGPrinter if self.config.output_format == "vcg" else DotPrinter
-        writer.DiagramWriter(self.config, printer_class).write(diadefs)
+        writer.DiagramWriter(self.config).write(diadefs)
         return 0
 
 

pylint/pyreverse/plantuml_printer.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2021 Andreas Finkler <andi.finkler@gmail.com>
+
+# 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.
+"""
+from typing import Dict, Optional
+
+from pylint.pyreverse.printer import EdgeType, Layout, NodeProperties, NodeType, Printer
+from pylint.pyreverse.utils import get_annotation_label
+
+
+class PlantUmlPrinter(Printer):
+    """Printer for PlantUML diagrams"""
+
+    DEFAULT_COLOR = "black"
+
+    NODES: Dict[NodeType, str] = {
+        NodeType.CLASS: "class",
+        NodeType.INTERFACE: "class",
+        NodeType.PACKAGE: "package",
+    }
+    ARROWS: Dict[EdgeType, str] = {
+        EdgeType.INHERITS: "--|>",
+        EdgeType.IMPLEMENTS: "..|>",
+        EdgeType.ASSOCIATION: "--*",
+        EdgeType.USES: "-->",
+    }
+
+    def _open_graph(self) -> None:
+        """Emit the header lines"""
+        self.emit("@startuml " + self.title)
+        if not self.use_automatic_namespace:
+            self.emit("set namespaceSeparator none")
+        if self.layout:
+            if self.layout is Layout.LEFT_TO_RIGHT:
+                self.emit("left to right direction")
+            elif self.layout is Layout.TOP_TO_BOTTOM:
+                self.emit("top to bottom direction")
+            else:
+                raise ValueError(
+                    f"Unsupported layout {self.layout}. PlantUmlPrinter only supports left to right and top to bottom layout."
+                )
+
+    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)
+        stereotype = " << interface >>" if type_ is NodeType.INTERFACE else ""
+        nodetype = self.NODES[type_]
+        if properties.color and properties.color != self.DEFAULT_COLOR:
+            color = f" #{properties.color}"
+        else:
+            color = ""
+        body = ""
+        if properties.attrs:
+            body += "\n".join(properties.attrs)
+        if properties.methods:
+            body += "\n"
+            for func in properties.methods:
+                args = self._get_method_arguments(func)
+                body += f"\n{func.name}({', '.join(args)})"
+                if func.returns:
+                    body += " -> " + get_annotation_label(func.returns)
+        label = properties.label if properties.label is not None else name
+        if properties.fontcolor and properties.fontcolor != self.DEFAULT_COLOR:
+            label = f"<color:{properties.fontcolor}>{label}</color>"
+        self.emit(f'{nodetype} "{label}" as {name}{stereotype}{color} {{\n{body}\n}}')
+
+    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."""
+        edge = f"{from_node} {self.ARROWS[type_]} {to_node}"
+        if label:
+            edge += f" : {label}"
+        self.emit(edge)
+
+    def _close_graph(self) -> None:
+        """Emit the lines needed to properly close the graph."""
+        self.emit("@enduml")

pylint/pyreverse/printer.py

@@ -12,6 +12,8 @@
 
 import astroid
 
+from pylint.pyreverse.utils import get_annotation_label
+
 
 class NodeType(Enum):
     CLASS = "class"
@@ -84,6 +86,28 @@ def emit_edge(
     ) -> None:
         """Create an edge from one node to another to display relationships."""
 
+    @staticmethod
+    def _get_method_arguments(method: astroid.FunctionDef) -> List[str]:
+        if method.args.args:
+            arguments: List[astroid.AssignName] = [
+                arg for arg in method.args.args if arg.name != "self"
+            ]
+        else:
+            arguments = []
+
+        annotations = dict(zip(arguments, method.args.annotations[1:]))
+        for arg in arguments:
+            annotation_label = ""
+            ann = annotations.get(arg)
+            if ann:
+                annotation_label = get_annotation_label(ann)
+            annotations[arg] = annotation_label
+
+        return [
+            f"{arg.name}: {ann}" if ann else f"{arg.name}"
+            for arg, ann in annotations.items()
+        ]
+
     def generate(self, outputfile: str) -> None:
         """Generate and save the final outputfile."""
         self._close_graph()

pylint/pyreverse/printer_factory.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2021 Andreas Finkler <andi.finkler@gmail.com>
+
+# 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
+
+from typing import Type
+
+from pylint.pyreverse.dot_printer import DotPrinter
+from pylint.pyreverse.plantuml_printer import PlantUmlPrinter
+from pylint.pyreverse.printer import Printer
+from pylint.pyreverse.vcg_printer import VCGPrinter
+
+filetype_to_printer = {
+    "vcg": VCGPrinter,
+    "plantuml": PlantUmlPrinter,
+    "puml": PlantUmlPrinter,
+    "dot": DotPrinter,
+}
+
+
+def get_printer_for_filetype(filetype: str) -> Type[Printer]:
+    return filetype_to_printer.get(filetype, DotPrinter)

pylint/pyreverse/writer.py

@@ -26,15 +26,16 @@
     PackageEntity,
 )
 from pylint.pyreverse.printer import EdgeType, NodeProperties, NodeType
+from pylint.pyreverse.printer_factory import get_printer_for_filetype
 from pylint.pyreverse.utils import is_exception
 
 
 class DiagramWriter:
     """base class for writing project diagrams"""
 
-    def __init__(self, config, printer_class):
+    def __init__(self, config):
         self.config = config
-        self.printer_class = printer_class
+        self.printer_class = get_printer_for_filetype(self.config.output_format)
         self.printer = None  # defined in set_printer
         self.file_name = ""  # defined in set_printer
 

tests/data/suppliermodule_test.py

@@ -7,6 +7,16 @@ def get_value(self):
     def set_value(self, value):
         raise NotImplementedError
 
+class CustomException(Exception): pass
+
 class DoNothing: pass
 
 class DoNothing2: pass
+
+class DoSomething:
+    def __init__(self, a_string: str, optional_int: int = None):
+        self.my_string = a_string
+        self.my_int = optional_int
+
+    def do_it(self, new_int: int) -> int:
+        return self.my_int + new_int

tests/pyreverse/conftest.py

@@ -54,6 +54,13 @@ def vcg_config() -> PyreverseConfig:
     )
 
 
+@pytest.fixture()
+def puml_config() -> PyreverseConfig:
+    return PyreverseConfig(
+        output_format="puml",
+    )
+
+
 @pytest.fixture(scope="session")
 def get_project() -> Callable:
     def _get_project(module: str, name: Optional[str] = "No Name") -> Project:

tests/pyreverse/data/classes_No_Name.dot

@@ -2,8 +2,10 @@ digraph "classes_No_Name" {
 rankdir=BT
 charset="utf-8"
 "data.clientmodule_test.Ancestor" [color="black", fontcolor="black", label="{Ancestor|attr : str\lcls_member\l|get_value()\lset_value(value)\l}", shape="record", style="solid"];
+"data.suppliermodule_test.CustomException" [color="black", fontcolor="red", label="{CustomException|\l|}", shape="record", style="solid"];
 "data.suppliermodule_test.DoNothing" [color="black", fontcolor="black", label="{DoNothing|\l|}", shape="record", style="solid"];
 "data.suppliermodule_test.DoNothing2" [color="black", fontcolor="black", label="{DoNothing2|\l|}", shape="record", style="solid"];
+"data.suppliermodule_test.DoSomething" [color="black", fontcolor="black", label="{DoSomething|my_int : Optional[int]\lmy_string : str\l|do_it(new_int: int): int\l}", shape="record", style="solid"];
 "data.suppliermodule_test.Interface" [color="black", fontcolor="black", label="{Interface|\l|get_value()\lset_value(value)\l}", shape="record", style="solid"];
 "data.clientmodule_test.Specialization" [color="black", fontcolor="black", label="{Specialization|TYPE : str\lrelation\lrelation2\ltop : str\l|}", shape="record", style="solid"];
 "data.clientmodule_test.Specialization" -> "data.clientmodule_test.Ancestor" [arrowhead="empty", arrowtail="none"];

tests/pyreverse/data/classes_No_Name.puml

@@ -0,0 +1,42 @@
+@startuml classes_No_Name
+set namespaceSeparator none
+class "Ancestor" as data.clientmodule_test.Ancestor {
+attr : str
+cls_member
+
+get_value()
+set_value(value)
+}
+class "<color:red>CustomException</color>" as data.suppliermodule_test.CustomException {
+
+}
+class "DoNothing" as data.suppliermodule_test.DoNothing {
+
+}
+class "DoNothing2" as data.suppliermodule_test.DoNothing2 {
+
+}
+class "DoSomething" as data.suppliermodule_test.DoSomething {
+my_int : Optional[int]
+my_string : str
+
+do_it(new_int: int) -> int
+}
+class "Interface" as data.suppliermodule_test.Interface {
+
+
+get_value()
+set_value(value)
+}
+class "Specialization" as data.clientmodule_test.Specialization {
+TYPE : str
+relation
+relation2
+top : str
+}
+data.clientmodule_test.Specialization --|> data.clientmodule_test.Ancestor
+data.clientmodule_test.Ancestor ..|> data.suppliermodule_test.Interface
+data.suppliermodule_test.DoNothing --* data.clientmodule_test.Ancestor : cls_member
+data.suppliermodule_test.DoNothing --* data.clientmodule_test.Specialization : relation
+data.suppliermodule_test.DoNothing2 --* data.clientmodule_test.Specialization : relation2
+@enduml

tests/pyreverse/data/classes_No_Name.vcg

@@ -6,12 +6,18 @@ graph:{
   manhattan_edges:yes
   node: {title:"data.clientmodule_test.Ancestor"  label:"\fbAncestor\fn\n\f____________\n\f08attr : str\n\f08cls_member\n\f____________\n\f10get_value()\n\f10set_value()"
   shape:box
+}
+  node: {title:"data.suppliermodule_test.CustomException"  label:"\fb09CustomException\fn\n\f_________________"
+  shape:box
 }
   node: {title:"data.suppliermodule_test.DoNothing"  label:"\fbDoNothing\fn\n\f___________"
   shape:box
 }
   node: {title:"data.suppliermodule_test.DoNothing2"  label:"\fbDoNothing2\fn\n\f____________"
   shape:box
+}
+  node: {title:"data.suppliermodule_test.DoSomething"  label:"\fbDoSomething\fn\n\f________________________\n\f08my_int : Optional[int]\n\f08my_string : str\n\f________________________\n\f10do_it()"
+  shape:box
 }
   node: {title:"data.suppliermodule_test.Interface"  label:"\fbInterface\fn\n\f___________\n\f10get_value()\n\f10set_value()"
   shape:box

tests/pyreverse/data/packages_No_Name.puml

@@ -0,0 +1,13 @@
+@startuml packages_No_Name
+set namespaceSeparator none
+package "data" as data {
+
+}
+package "data.clientmodule_test" as data.clientmodule_test {
+
+}
+package "data.suppliermodule_test" as data.suppliermodule_test {
+
+}
+data.clientmodule_test --> data.suppliermodule_test
+@enduml

tests/pyreverse/test_diadefs.py

@@ -141,8 +141,10 @@ def test_known_values1(HANDLER, PROJECT):
     classes = _process_classes(cd.objects)
     assert classes == [
         (True, "Ancestor"),
+        (True, "CustomException"),
         (True, "DoNothing"),
         (True, "DoNothing2"),
+        (True, "DoSomething"),
         (True, "Interface"),
         (True, "Specialization"),
     ]