``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

18 files changed, 310 insertions, 55 deletions

diff --git a/ChangeLog b/ChangeLog
index 02ad13eea..003c2f094 100644
--- a/ChangeLog
+++ b/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.
diff --git a/doc/whatsnew/2.10.rst b/doc/whatsnew/2.10.rst
index 425135194..37b2ca854 100644
--- a/doc/whatsnew/2.10.rst
+++ b/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.
diff --git a/pylint/pyreverse/dot_printer.py b/pylint/pyreverse/dot_printer.py
index bcd58984b..4b735bb3f 100644
--- a/pylint/pyreverse/dot_printer.py
+++ b/pylint/pyreverse/dot_printer.py
@@ -80,9 +80,8 @@ class DotPrinter(Printer):
             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 @@ class DotPrinter(Printer):
         # 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
 
diff --git a/pylint/pyreverse/main.py b/pylint/pyreverse/main.py
index 323ba0788..163eb4ecb 100644
--- a/pylint/pyreverse/main.py
+++ b/pylint/pyreverse/main.py
@@ -27,10 +27,8 @@ from typing import Iterable
 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 @@ class Run(ConfigurationMixIn):
             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
 
 
diff --git a/pylint/pyreverse/plantuml_printer.py b/pylint/pyreverse/plantuml_printer.py
new file mode 100644
index 000000000..4b362db1a
--- /dev/null
+++ b/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")
diff --git a/pylint/pyreverse/printer.py b/pylint/pyreverse/printer.py
index c91d75560..52a51a234 100644
--- a/pylint/pyreverse/printer.py
+++ b/pylint/pyreverse/printer.py
@@ -12,6 +12,8 @@ from typing import List, NamedTuple, Optional
 
 import astroid
 
+from pylint.pyreverse.utils import get_annotation_label
+
 
 class NodeType(Enum):
     CLASS = "class"
@@ -84,6 +86,28 @@ class Printer(ABC):
     ) -> 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()
diff --git a/pylint/pyreverse/printer_factory.py b/pylint/pyreverse/printer_factory.py
new file mode 100644
index 000000000..c52f24a4d
--- /dev/null
+++ b/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)
diff --git a/pylint/pyreverse/writer.py b/pylint/pyreverse/writer.py
index 4a4c51008..716d55d5d 100644
--- a/pylint/pyreverse/writer.py
+++ b/pylint/pyreverse/writer.py
@@ -26,15 +26,16 @@ from pylint.pyreverse.diagrams import (
     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
 
diff --git a/tests/data/suppliermodule_test.py b/tests/data/suppliermodule_test.py
index 6af30fa08..add637d4a 100644
--- a/tests/data/suppliermodule_test.py
+++ b/tests/data/suppliermodule_test.py
@@ -7,6 +7,16 @@ class Interface:
     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
diff --git a/tests/pyreverse/conftest.py b/tests/pyreverse/conftest.py
index f079f7105..b7ccfce06 100644
--- a/tests/pyreverse/conftest.py
+++ b/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:
diff --git a/tests/pyreverse/data/classes_No_Name.dot b/tests/pyreverse/data/classes_No_Name.dot
index c0141fd2a..f3e7b9625 100644
--- a/tests/pyreverse/data/classes_No_Name.dot
+++ b/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"];
diff --git a/tests/pyreverse/data/classes_No_Name.puml b/tests/pyreverse/data/classes_No_Name.puml
new file mode 100644
index 000000000..c833d7fef
--- /dev/null
+++ b/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
diff --git a/tests/pyreverse/data/classes_No_Name.vcg b/tests/pyreverse/data/classes_No_Name.vcg
index 1bf83e338..89697ce84 100644
--- a/tests/pyreverse/data/classes_No_Name.vcg
+++ b/tests/pyreverse/data/classes_No_Name.vcg
@@ -7,12 +7,18 @@ graph:{
   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
 }
diff --git a/tests/pyreverse/data/packages_No_Name.puml b/tests/pyreverse/data/packages_No_Name.puml
new file mode 100644
index 000000000..7089df0ea
--- /dev/null
+++ b/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
diff --git a/tests/pyreverse/test_diadefs.py b/tests/pyreverse/test_diadefs.py
index 3f08c1f3d..e38d1dc1f 100644
--- a/tests/pyreverse/test_diadefs.py
+++ b/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"),
     ]
diff --git a/tests/pyreverse/test_printer.py b/tests/pyreverse/test_printer.py
index 44dcefc76..703d9dfc9 100644
--- a/tests/pyreverse/test_printer.py
+++ b/tests/pyreverse/test_printer.py
@@ -8,7 +8,8 @@ from typing import Type
 import pytest
 
 from pylint.pyreverse.dot_printer import DotPrinter
-from pylint.pyreverse.printer import Layout, Printer
+from pylint.pyreverse.plantuml_printer import PlantUmlPrinter
+from pylint.pyreverse.printer import Layout, NodeType, Printer
 from pylint.pyreverse.vcg_printer import VCGPrinter
 
 
@@ -23,6 +24,8 @@ from pylint.pyreverse.vcg_printer import VCGPrinter
         (Layout.BOTTOM_TO_TOP, VCGPrinter, "orientation:bottom_to_top", -1),
         (Layout.LEFT_TO_RIGHT, VCGPrinter, "orientation:left_to_right", -1),
         (Layout.RIGHT_TO_LEFT, VCGPrinter, "orientation:right_to_left", -1),
+        (Layout.TOP_TO_BOTTOM, PlantUmlPrinter, "top to bottom direction", -1),
+        (Layout.LEFT_TO_RIGHT, PlantUmlPrinter, "left to right direction", -1),
     ],
 )
 def test_explicit_layout(
@@ -30,3 +33,20 @@ def test_explicit_layout(
 ) -> None:
     printer = printer_class(title="unittest", layout=layout)
     assert printer.lines[line_index].strip() == expected_content
+
+
+@pytest.mark.parametrize(
+    "layout, printer_class",
+    [(Layout.BOTTOM_TO_TOP, PlantUmlPrinter), (Layout.RIGHT_TO_LEFT, PlantUmlPrinter)],
+)
+def test_unsupported_layout(layout: Layout, printer_class: Type[Printer]):
+    with pytest.raises(ValueError):
+        printer_class(title="unittest", layout=layout)
+
+
+class TestPlantUmlPrinter:
+    printer = PlantUmlPrinter(title="unittest", layout=Layout.TOP_TO_BOTTOM)
+
+    def test_node_without_properties(self):
+        self.printer.emit_node(name="test", type_=NodeType.CLASS)
+        assert self.printer.lines[-1] == 'class "test" as test {\n\n}\n'
diff --git a/tests/pyreverse/test_printer_factory.py b/tests/pyreverse/test_printer_factory.py
new file mode 100644
index 000000000..f2ad7106a
--- /dev/null
+++ b/tests/pyreverse/test_printer_factory.py
@@ -0,0 +1,29 @@
+# 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
+
+"""
+Unit tests for pylint.pyreverse.printer_factory
+"""
+
+import pytest
+
+from pylint.pyreverse import printer_factory
+from pylint.pyreverse.dot_printer import DotPrinter
+from pylint.pyreverse.plantuml_printer import PlantUmlPrinter
+from pylint.pyreverse.vcg_printer import VCGPrinter
+
+
+@pytest.mark.parametrize(
+    "filetype, expected_printer_class",
+    [
+        ("vcg", VCGPrinter),
+        ("dot", DotPrinter),
+        ("puml", PlantUmlPrinter),
+        ("plantuml", PlantUmlPrinter),
+        ("png", DotPrinter),
+    ],
+)
+def test_get_printer_for_filetype(filetype, expected_printer_class):
+    assert printer_factory.get_printer_for_filetype(filetype) == expected_printer_class
diff --git a/tests/pyreverse/test_writer.py b/tests/pyreverse/test_writer.py
index 3ed1d1e68..167efb614 100644
--- a/tests/pyreverse/test_writer.py
+++ b/tests/pyreverse/test_writer.py
@@ -28,9 +28,7 @@ from difflib import unified_diff
 import pytest
 
 from pylint.pyreverse.diadefslib import DefaultDiadefGenerator, DiadefsHandler
-from pylint.pyreverse.dot_printer import DotPrinter
 from pylint.pyreverse.inspector import Linker
-from pylint.pyreverse.vcg_printer import VCGPrinter
 from pylint.pyreverse.writer import DiagramWriter
 
 _DEFAULTS = {
@@ -74,22 +72,28 @@ def _file_lines(path):
 
 DOT_FILES = ["packages_No_Name.dot", "classes_No_Name.dot"]
 VCG_FILES = ["packages_No_Name.vcg", "classes_No_Name.vcg"]
+PUML_FILES = ["packages_No_Name.puml", "classes_No_Name.puml"]
 
 
 @pytest.fixture()
 def setup_dot(default_config, get_project):
-    config = default_config
-    writer = DiagramWriter(config, printer_class=DotPrinter)
+    writer = DiagramWriter(default_config)
     project = get_project(os.path.join(os.path.dirname(__file__), "..", "data"))
-    yield from _setup(project, config, writer)
+    yield from _setup(project, default_config, writer)
 
 
 @pytest.fixture()
 def setup_vcg(vcg_config, get_project):
-    config = vcg_config
-    writer = DiagramWriter(config, printer_class=VCGPrinter)
+    writer = DiagramWriter(vcg_config)
     project = get_project(os.path.join(os.path.dirname(__file__), "..", "data"))
-    yield from _setup(project, config, writer)
+    yield from _setup(project, vcg_config, writer)
+
+
+@pytest.fixture()
+def setup_puml(puml_config, get_project):
+    writer = DiagramWriter(puml_config)
+    project = get_project(os.path.join(os.path.dirname(__file__), "..", "data"))
+    yield from _setup(project, puml_config, writer)
 
 
 def _setup(project, config, writer):
@@ -100,7 +104,7 @@ def _setup(project, config, writer):
         diagram.extract_relationships()
     writer.write(dd)
     yield
-    for fname in DOT_FILES + VCG_FILES:
+    for fname in DOT_FILES + VCG_FILES + PUML_FILES:
         try:
             os.remove(fname)
         except FileNotFoundError:
@@ -110,22 +114,22 @@ def _setup(project, config, writer):
 @pytest.mark.usefixtures("setup_dot")
 @pytest.mark.parametrize("generated_file", DOT_FILES)
 def test_dot_files(generated_file):
-    expected_file = os.path.join(os.path.dirname(__file__), "data", generated_file)
-    generated = _file_lines(generated_file)
-    expected = _file_lines(expected_file)
-    generated = "\n".join(generated)
-    expected = "\n".join(expected)
-    files = f"\n *** expected : {expected_file}, generated : {generated_file} \n"
-    diff = "\n".join(
-        line for line in unified_diff(expected.splitlines(), generated.splitlines())
-    )
-    assert expected == generated, f"{files}{diff}"
-    os.remove(generated_file)
+    _assert_files_are_equal(generated_file)
 
 
 @pytest.mark.usefixtures("setup_vcg")
 @pytest.mark.parametrize("generated_file", VCG_FILES)
 def test_vcg_files(generated_file):
+    _assert_files_are_equal(generated_file)
+
+
+@pytest.mark.usefixtures("setup_puml")
+@pytest.mark.parametrize("generated_file", PUML_FILES)
+def test_puml_files(generated_file):
+    _assert_files_are_equal(generated_file)
+
+
+def _assert_files_are_equal(generated_file):
     expected_file = os.path.join(os.path.dirname(__file__), "data", generated_file)
     generated = _file_lines(generated_file)
     expected = _file_lines(expected_file)
@@ -136,4 +140,3 @@ def test_vcg_files(generated_file):
         line for line in unified_diff(expected.splitlines(), generated.splitlines())
     )
     assert expected == generated, f"{files}{diff}"
-    os.remove(generated_file)