``pyreverse``: Add option for colored output (#4850)

* Add option to produce colored output from ``pyreverse``

* Use indentation in PlantUML diagrams

Co-authored-by: Pierre Sassoulas <pierre.sassoulas@gmail.com>

16 files changed, 281 insertions, 86 deletions

diff --git a/ChangeLog b/ChangeLog
index 003c2f094..1c13c66d1 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -12,7 +12,11 @@ Release date: TBA
 ..
   Put bug fixes that should not wait for a new minor version here
 
-+ pyreverse: add output in PlantUML format.
+* pyreverse: add option to produce colored output.
+
+  Closes #4488
+
+* pyreverse: add output in PlantUML format.
 
   Closes #4498
 
diff --git a/doc/whatsnew/2.10.rst b/doc/whatsnew/2.10.rst
index 37b2ca854..e03e4d04a 100644
--- a/doc/whatsnew/2.10.rst
+++ b/doc/whatsnew/2.10.rst
@@ -44,6 +44,8 @@ Extensions
 Other Changes
 =============
 
+* pyreverse now permit to produce colored generated diagram by using the ``colorized`` option.
+
 * Pyreverse - add output in PlantUML format
 
 * pylint does not crash with a traceback anymore when a file is problematic. It
diff --git a/pylint/pyreverse/dot_printer.py b/pylint/pyreverse/dot_printer.py
index 4b735bb3f..eeef13766 100644
--- a/pylint/pyreverse/dot_printer.py
+++ b/pylint/pyreverse/dot_printer.py
@@ -35,6 +35,8 @@ ARROWS: Dict[EdgeType, Dict] = {
 
 
 class DotPrinter(Printer):
+    DEFAULT_COLOR = "black"
+
     def __init__(
         self,
         title: str,
@@ -43,7 +45,6 @@ class DotPrinter(Printer):
     ):
         layout = layout or Layout.BOTTOM_TO_TOP
         self.charset = "utf-8"
-        self.node_style = "solid"
         super().__init__(title, layout, use_automatic_namespace)
 
     def _open_graph(self) -> None:
@@ -67,7 +68,8 @@ class DotPrinter(Printer):
         if properties is None:
             properties = NodeProperties(label=name)
         shape = SHAPES[type_]
-        color = properties.color if properties.color is not None else "black"
+        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)
         if label:
             label_part = f', label="{label}"'
@@ -77,7 +79,7 @@ class DotPrinter(Printer):
             f', fontcolor="{properties.fontcolor}"' if properties.fontcolor else ""
         )
         self.emit(
-            f'"{name}" [color="{color}"{fontcolor_part}{label_part}, shape="{shape}", style="{self.node_style}"];'
+            f'"{name}" [color="{color}"{fontcolor_part}{label_part}, shape="{shape}", style="{style}"];'
         )
 
     def _build_label_for_node(
diff --git a/pylint/pyreverse/main.py b/pylint/pyreverse/main.py
index 163eb4ecb..1057e9ce5 100644
--- a/pylint/pyreverse/main.py
+++ b/pylint/pyreverse/main.py
@@ -11,6 +11,7 @@
 # Copyright (c) 2020 hippo91 <guillaume.peillex@gmail.com>
 # Copyright (c) 2021 Marc Mueller <30130371+cdce8p@users.noreply.github.com>
 # Copyright (c) 2021 Mark Byrne <31762852+mbyrnepr2@users.noreply.github.com>
+# 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
@@ -124,8 +125,7 @@ OPTIONS = (
             short="k",
             action="store_true",
             default=False,
-            help="don't show attributes and methods in the class boxes; \
-this disables -f values",
+            help="don't show attributes and methods in the class boxes; this disables -f values",
         ),
     ),
     (
@@ -140,36 +140,55 @@ this disables -f values",
         ),
     ),
     (
+        "colorized",
+        dict(
+            dest="colorized",
+            action="store_true",
+            default=False,
+            help="Use colored output. Classes/modules of the same package get the same color.",
+        ),
+    ),
+    (
+        "max-color-depth",
+        dict(
+            dest="max_color_depth",
+            action="store",
+            default=2,
+            metavar="<depth>",
+            type="int",
+            help="Use separate colors up to package depth of <depth>",
+        ),
+    ),
+    (
         "ignore",
-        {
-            "type": "csv",
-            "metavar": "<file[,file...]>",
-            "dest": "ignore_list",
-            "default": ("CVS",),
-            "help": "Files or directories to be skipped. They "
-            "should be base names, not paths.",
-        },
+        dict(
+            type="csv",
+            metavar="<file[,file...]>",
+            dest="ignore_list",
+            default=("CVS",),
+            help="Files or directories to be skipped. They should be base names, not paths.",
+        ),
     ),
     (
         "project",
-        {
-            "default": "",
-            "type": "string",
-            "short": "p",
-            "metavar": "<project name>",
-            "help": "set the project name.",
-        },
+        dict(
+            default="",
+            type="string",
+            short="p",
+            metavar="<project name>",
+            help="set the project name.",
+        ),
     ),
     (
         "output-directory",
-        {
-            "default": "",
-            "type": "string",
-            "short": "d",
-            "action": "store",
-            "metavar": "<output_directory>",
-            "help": "set the output directory path.",
-        },
+        dict(
+            default="",
+            type="string",
+            short="d",
+            action="store",
+            metavar="<output_directory>",
+            help="set the output directory path.",
+        ),
     ),
 )
 
diff --git a/pylint/pyreverse/plantuml_printer.py b/pylint/pyreverse/plantuml_printer.py
index 4b362db1a..c73cddd5e 100644
--- a/pylint/pyreverse/plantuml_printer.py
+++ b/pylint/pyreverse/plantuml_printer.py
@@ -59,20 +59,25 @@ class PlantUmlPrinter(Printer):
             color = f" #{properties.color}"
         else:
             color = ""
-        body = ""
+        body = []
         if properties.attrs:
-            body += "\n".join(properties.attrs)
+            body.extend(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)})"
+                line = f"{func.name}({', '.join(args)})"
                 if func.returns:
-                    body += " -> " + get_annotation_label(func.returns)
+                    line += " -> " + get_annotation_label(func.returns)
+                body.append(line)
         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}}')
+        self.emit(f'{nodetype} "{label}" as {name}{stereotype}{color} {{')
+        self._inc_indent()
+        for line in body:
+            self.emit(line)
+        self._dec_indent()
+        self.emit("}")
 
     def emit_edge(
         self,
diff --git a/pylint/pyreverse/printer.py b/pylint/pyreverse/printer.py
index 52a51a234..fe9a7b7a5 100644
--- a/pylint/pyreverse/printer.py
+++ b/pylint/pyreverse/printer.py
@@ -56,8 +56,17 @@ class Printer(ABC):
         self.layout = layout
         self.use_automatic_namespace = use_automatic_namespace
         self.lines: List[str] = []
+        self._indent = ""
         self._open_graph()
 
+    def _inc_indent(self):
+        """increment indentation"""
+        self._indent += "  "
+
+    def _dec_indent(self):
+        """decrement indentation"""
+        self._indent = self._indent[:-2]
+
     @abstractmethod
     def _open_graph(self) -> None:
         """Emit the header lines, i.e. all boilerplate code that defines things like layout etc."""
@@ -65,7 +74,7 @@ class Printer(ABC):
     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)
+        self.lines.append(self._indent + line)
 
     @abstractmethod
     def emit_node(
diff --git a/pylint/pyreverse/vcg_printer.py b/pylint/pyreverse/vcg_printer.py
index 62559d675..8868443d2 100644
--- a/pylint/pyreverse/vcg_printer.py
+++ b/pylint/pyreverse/vcg_printer.py
@@ -185,18 +185,9 @@ ORIENTATION: Dict[Layout, str] = {
 
 
 class VCGPrinter(Printer):
-    def __init__(
-        self,
-        title: str,
-        layout: Optional[Layout] = None,
-        use_automatic_namespace: Optional[bool] = None,
-    ):
-        self._indent = ""
-        super().__init__(title, layout, use_automatic_namespace)
-
     def _open_graph(self) -> None:
         """Emit the header lines"""
-        self.emit(f"{self._indent}graph:{{\n")
+        self.emit("graph:{\n")
         self._inc_indent()
         self._write_attributes(
             GRAPH_ATTRS,
@@ -212,7 +203,7 @@ class VCGPrinter(Printer):
     def _close_graph(self) -> None:
         """Emit the lines needed to properly close the graph."""
         self._dec_indent()
-        self.emit(f"{self._indent}}}")
+        self.emit("}")
 
     def emit_node(
         self,
@@ -225,7 +216,7 @@ class VCGPrinter(Printer):
             properties = NodeProperties(label=name)
         elif properties.label is None:
             properties.label = name
-        self.emit(f'{self._indent}node: {{title:"{name}"', force_newline=False)
+        self.emit(f'node: {{title:"{name}"', force_newline=False)
         self._write_attributes(
             NODE_ATTRS,
             label=self._build_label_for_node(properties),
@@ -264,7 +255,7 @@ class VCGPrinter(Printer):
     ) -> None:
         """Create an edge from one node to another to display relationships."""
         self.emit(
-            f'{self._indent}edge: {{sourcename:"{from_node}" targetname:"{to_node}"',
+            f'edge: {{sourcename:"{from_node}" targetname:"{to_node}"',
             force_newline=False,
         )
         attributes = ARROWS[type_]
@@ -287,20 +278,12 @@ class VCGPrinter(Printer):
                 ) from e
 
             if not _type:
-                self.emit(f'{self._indent}{key}:"{value}"\n')
+                self.emit(f'{key}:"{value}"\n')
             elif _type == 1:
-                self.emit(f"{self._indent}{key}:{int(value)}\n")
+                self.emit(f"{key}:{int(value)}\n")
             elif value in _type:
-                self.emit(f"{self._indent}{key}:{value}\n")
+                self.emit(f"{key}:{value}\n")
             else:
                 raise Exception(
                     f"value {value} isn't correct for attribute {key} correct values are {type}"
                 )
-
-    def _inc_indent(self):
-        """increment indentation"""
-        self._indent += "  "
-
-    def _dec_indent(self):
-        """decrement indentation"""
-        self._indent = self._indent[:-2]
diff --git a/pylint/pyreverse/writer.py b/pylint/pyreverse/writer.py
index 716d55d5d..9f0660545 100644
--- a/pylint/pyreverse/writer.py
+++ b/pylint/pyreverse/writer.py
@@ -17,11 +17,16 @@
 
 """Utilities for creating VCG and Dot diagrams"""
 
+import itertools
 import os
 
+import astroid
+from astroid import modutils
+
 from pylint.pyreverse.diagrams import (
     ClassDiagram,
     ClassEntity,
+    DiagramEntity,
     PackageDiagram,
     PackageEntity,
 )
@@ -38,6 +43,29 @@ class DiagramWriter:
         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
+        self.depth = self.config.max_color_depth
+        self.available_colors = itertools.cycle(
+            [
+                "aliceblue",
+                "antiquewhite",
+                "aquamarine",
+                "burlywood",
+                "cadetblue",
+                "chartreuse",
+                "chocolate",
+                "coral",
+                "cornflowerblue",
+                "cyan",
+                "darkgoldenrod",
+                "darkseagreen",
+                "dodgerblue",
+                "forestgreen",
+                "gold",
+                "hotpink",
+                "mediumspringgreen",
+            ]
+        )
+        self.used_colors = {}
 
     def write(self, diadefs):
         """write files for <project> according to <diadefs>"""
@@ -108,12 +136,11 @@ class DiagramWriter:
         self.printer = self.printer_class(basename)
         self.file_name = file_name
 
-    @staticmethod
-    def get_package_properties(obj: PackageEntity) -> NodeProperties:
+    def get_package_properties(self, obj: PackageEntity) -> NodeProperties:
         """get label and shape for packages."""
         return NodeProperties(
             label=obj.title,
-            color="black",
+            color=self.get_shape_color(obj) if self.config.colorized else "black",
         )
 
     def get_class_properties(self, obj: ClassEntity) -> NodeProperties:
@@ -123,10 +150,26 @@ class DiagramWriter:
             attrs=obj.attrs if not self.config.only_classnames else None,
             methods=obj.methods if not self.config.only_classnames else None,
             fontcolor="red" if is_exception(obj.node) else "black",
-            color="black",
+            color=self.get_shape_color(obj) if self.config.colorized else "black",
         )
         return properties
 
+    def get_shape_color(self, obj: DiagramEntity) -> str:
+        """get shape color"""
+        qualified_name = obj.node.qname()
+        if modutils.is_standard_module(qualified_name.split(".", maxsplit=1)[0]):
+            return "grey"
+        if isinstance(obj.node, astroid.ClassDef):
+            package = qualified_name.rsplit(".", maxsplit=2)[0]
+        elif obj.node.package:
+            package = qualified_name
+        else:
+            package = qualified_name.rsplit(".", maxsplit=1)[0]
+        base_name = ".".join(package.split(".", self.depth)[: self.depth])
+        if base_name not in self.used_colors:
+            self.used_colors[base_name] = next(self.available_colors)
+        return self.used_colors[base_name]
+
     def save(self) -> None:
         """write to disk"""
         self.printer.generate(self.file_name)
diff --git a/tests/pyreverse/conftest.py b/tests/pyreverse/conftest.py
index b7ccfce06..982980101 100644
--- a/tests/pyreverse/conftest.py
+++ b/tests/pyreverse/conftest.py
@@ -23,6 +23,8 @@ class PyreverseConfig:  # pylint: disable=too-many-instance-attributes, too-many
         module_names: Optional[bool] = None,
         only_classnames: bool = False,
         output_format: str = "dot",
+        colorized: bool = False,
+        max_color_depth: int = 2,
         ignore_list: Tuple = tuple(),
         project: str = "",
         output_directory: str = "",
@@ -37,6 +39,8 @@ class PyreverseConfig:  # pylint: disable=too-many-instance-attributes, too-many
         self.module_names = module_names
         self.only_classnames = only_classnames
         self.output_format = output_format
+        self.colorized = colorized
+        self.max_color_depth = max_color_depth
         self.ignore_list = ignore_list
         self.project = project
         self.output_directory = output_directory
@@ -48,6 +52,14 @@ def default_config() -> PyreverseConfig:
 
 
 @pytest.fixture()
+def colorized_dot_config() -> PyreverseConfig:
+    return PyreverseConfig(
+        output_format="dot",
+        colorized=True,
+    )
+
+
+@pytest.fixture()
 def vcg_config() -> PyreverseConfig:
     return PyreverseConfig(
         output_format="vcg",
@@ -61,6 +73,14 @@ def puml_config() -> PyreverseConfig:
     )
 
 
+@pytest.fixture()
+def colorized_puml_config() -> PyreverseConfig:
+    return PyreverseConfig(
+        output_format="puml",
+        colorized=True,
+    )
+
+
 @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.puml b/tests/pyreverse/data/classes_No_Name.puml
index c833d7fef..5acbe93d0 100644
--- a/tests/pyreverse/data/classes_No_Name.puml
+++ b/tests/pyreverse/data/classes_No_Name.puml
@@ -1,38 +1,31 @@
 @startuml classes_No_Name
 set namespaceSeparator none
 class "Ancestor" as data.clientmodule_test.Ancestor {
-attr : str
-cls_member
-
-get_value()
-set_value(value)
+  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
+  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)
+  get_value()
+  set_value(value)
 }
 class "Specialization" as data.clientmodule_test.Specialization {
-TYPE : str
-relation
-relation2
-top : str
+  TYPE : str
+  relation
+  relation2
+  top : str
 }
 data.clientmodule_test.Specialization --|> data.clientmodule_test.Ancestor
 data.clientmodule_test.Ancestor ..|> data.suppliermodule_test.Interface
diff --git a/tests/pyreverse/data/classes_colorized.dot b/tests/pyreverse/data/classes_colorized.dot
new file mode 100644
index 000000000..41c81d94f
--- /dev/null
+++ b/tests/pyreverse/data/classes_colorized.dot
@@ -0,0 +1,16 @@
+digraph "classes_colorized" {
+rankdir=BT
+charset="utf-8"
+"data.clientmodule_test.Ancestor" [color="aliceblue", fontcolor="black", label="{Ancestor|attr : str\lcls_member\l|get_value()\lset_value(value)\l}", shape="record", style="filled"];
+"data.suppliermodule_test.CustomException" [color="aliceblue", fontcolor="red", label="{CustomException|\l|}", shape="record", style="filled"];
+"data.suppliermodule_test.DoNothing" [color="aliceblue", fontcolor="black", label="{DoNothing|\l|}", shape="record", style="filled"];
+"data.suppliermodule_test.DoNothing2" [color="aliceblue", fontcolor="black", label="{DoNothing2|\l|}", shape="record", style="filled"];
+"data.suppliermodule_test.DoSomething" [color="aliceblue", fontcolor="black", label="{DoSomething|my_int : Optional[int]\lmy_string : str\l|do_it(new_int: int): int\l}", shape="record", style="filled"];
+"data.suppliermodule_test.Interface" [color="aliceblue", fontcolor="black", label="{Interface|\l|get_value()\lset_value(value)\l}", shape="record", style="filled"];
+"data.clientmodule_test.Specialization" [color="aliceblue", fontcolor="black", label="{Specialization|TYPE : str\lrelation\lrelation2\ltop : str\l|}", shape="record", style="filled"];
+"data.clientmodule_test.Specialization" -> "data.clientmodule_test.Ancestor" [arrowhead="empty", arrowtail="none"];
+"data.clientmodule_test.Ancestor" -> "data.suppliermodule_test.Interface" [arrowhead="empty", arrowtail="node", style="dashed"];
+"data.suppliermodule_test.DoNothing" -> "data.clientmodule_test.Ancestor" [arrowhead="diamond", arrowtail="none", fontcolor="green", label="cls_member", style="solid"];
+"data.suppliermodule_test.DoNothing" -> "data.clientmodule_test.Specialization" [arrowhead="diamond", arrowtail="none", fontcolor="green", label="relation", style="solid"];
+"data.suppliermodule_test.DoNothing2" -> "data.clientmodule_test.Specialization" [arrowhead="diamond", arrowtail="none", fontcolor="green", label="relation2", style="solid"];
+}
diff --git a/tests/pyreverse/data/classes_colorized.puml b/tests/pyreverse/data/classes_colorized.puml
new file mode 100644
index 000000000..cbbc3b02e
--- /dev/null
+++ b/tests/pyreverse/data/classes_colorized.puml
@@ -0,0 +1,35 @@
+@startuml classes_colorized
+set namespaceSeparator none
+class "Ancestor" as data.clientmodule_test.Ancestor #aliceblue {
+  attr : str
+  cls_member
+  get_value()
+  set_value(value)
+}
+class "<color:red>CustomException</color>" as data.suppliermodule_test.CustomException #aliceblue {
+}
+class "DoNothing" as data.suppliermodule_test.DoNothing #aliceblue {
+}
+class "DoNothing2" as data.suppliermodule_test.DoNothing2 #aliceblue {
+}
+class "DoSomething" as data.suppliermodule_test.DoSomething #aliceblue {
+  my_int : Optional[int]
+  my_string : str
+  do_it(new_int: int) -> int
+}
+class "Interface" as data.suppliermodule_test.Interface #aliceblue {
+  get_value()
+  set_value(value)
+}
+class "Specialization" as data.clientmodule_test.Specialization #aliceblue {
+  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/packages_colorized.dot b/tests/pyreverse/data/packages_colorized.dot
new file mode 100644
index 000000000..917fa86f5
--- /dev/null
+++ b/tests/pyreverse/data/packages_colorized.dot
@@ -0,0 +1,8 @@
+digraph "packages_colorized" {
+rankdir=BT
+charset="utf-8"
+"data" [color="aliceblue", label="data", shape="box", style="filled"];
+"data.clientmodule_test" [color="aliceblue", label="data.clientmodule_test", shape="box", style="filled"];
+"data.suppliermodule_test" [color="aliceblue", label="data.suppliermodule_test", shape="box", style="filled"];
+"data.clientmodule_test" -> "data.suppliermodule_test" [arrowhead="open", arrowtail="none"];
+}
diff --git a/tests/pyreverse/data/packages_colorized.puml b/tests/pyreverse/data/packages_colorized.puml
new file mode 100644
index 000000000..8d489e7aa
--- /dev/null
+++ b/tests/pyreverse/data/packages_colorized.puml
@@ -0,0 +1,13 @@
+@startuml packages_colorized
+set namespaceSeparator none
+package "data" as data #aliceblue {
+
+}
+package "data.clientmodule_test" as data.clientmodule_test #aliceblue {
+
+}
+package "data.suppliermodule_test" as data.suppliermodule_test #aliceblue {
+
+}
+data.clientmodule_test --> data.suppliermodule_test
+@enduml
diff --git a/tests/pyreverse/test_printer.py b/tests/pyreverse/test_printer.py
index 703d9dfc9..77d12f630 100644
--- a/tests/pyreverse/test_printer.py
+++ b/tests/pyreverse/test_printer.py
@@ -49,4 +49,4 @@ class TestPlantUmlPrinter:
 
     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'
+        assert self.printer.lines[-2:] == ['class "test" as test {\n', "}\n"]
diff --git a/tests/pyreverse/test_writer.py b/tests/pyreverse/test_writer.py
index 167efb614..3b92ab39a 100644
--- a/tests/pyreverse/test_writer.py
+++ b/tests/pyreverse/test_writer.py
@@ -24,6 +24,7 @@ Unit test for ``DiagramWriter``
 import codecs
 import os
 from difflib import unified_diff
+from unittest.mock import Mock
 
 import pytest
 
@@ -71,8 +72,10 @@ def _file_lines(path):
 
 
 DOT_FILES = ["packages_No_Name.dot", "classes_No_Name.dot"]
+COLORIZED_DOT_FILES = ["packages_colorized.dot", "classes_colorized.dot"]
 VCG_FILES = ["packages_No_Name.vcg", "classes_No_Name.vcg"]
 PUML_FILES = ["packages_No_Name.puml", "classes_No_Name.puml"]
+COLORIZED_PUML_FILES = ["packages_colorized.puml", "classes_colorized.puml"]
 
 
 @pytest.fixture()
@@ -83,6 +86,15 @@ def setup_dot(default_config, get_project):
 
 
 @pytest.fixture()
+def setup_colorized_dot(colorized_dot_config, get_project):
+    writer = DiagramWriter(colorized_dot_config)
+    project = get_project(
+        os.path.join(os.path.dirname(__file__), "..", "data"), name="colorized"
+    )
+    yield from _setup(project, colorized_dot_config, writer)
+
+
+@pytest.fixture()
 def setup_vcg(vcg_config, get_project):
     writer = DiagramWriter(vcg_config)
     project = get_project(os.path.join(os.path.dirname(__file__), "..", "data"))
@@ -96,6 +108,15 @@ def setup_puml(puml_config, get_project):
     yield from _setup(project, puml_config, writer)
 
 
+@pytest.fixture()
+def setup_colorized_puml(colorized_puml_config, get_project):
+    writer = DiagramWriter(colorized_puml_config)
+    project = get_project(
+        os.path.join(os.path.dirname(__file__), "..", "data"), name="colorized"
+    )
+    yield from _setup(project, colorized_puml_config, writer)
+
+
 def _setup(project, config, writer):
     linker = Linker(project)
     handler = DiadefsHandler(config)
@@ -104,7 +125,9 @@ def _setup(project, config, writer):
         diagram.extract_relationships()
     writer.write(dd)
     yield
-    for fname in DOT_FILES + VCG_FILES + PUML_FILES:
+    for fname in (
+        DOT_FILES + COLORIZED_DOT_FILES + VCG_FILES + PUML_FILES + COLORIZED_PUML_FILES
+    ):
         try:
             os.remove(fname)
         except FileNotFoundError:
@@ -117,6 +140,12 @@ def test_dot_files(generated_file):
     _assert_files_are_equal(generated_file)
 
 
+@pytest.mark.usefixtures("setup_colorized_dot")
+@pytest.mark.parametrize("generated_file", COLORIZED_DOT_FILES)
+def test_colorized_dot_files(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):
@@ -129,6 +158,12 @@ def test_puml_files(generated_file):
     _assert_files_are_equal(generated_file)
 
 
+@pytest.mark.usefixtures("setup_colorized_puml")
+@pytest.mark.parametrize("generated_file", COLORIZED_PUML_FILES)
+def test_colorized_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)
@@ -140,3 +175,11 @@ def _assert_files_are_equal(generated_file):
         line for line in unified_diff(expected.splitlines(), generated.splitlines())
     )
     assert expected == generated, f"{files}{diff}"
+
+
+def test_color_for_stdlib_module(default_config):
+    writer = DiagramWriter(default_config)
+    obj = Mock()
+    obj.node = Mock()
+    obj.node.qname.return_value = "collections"
+    assert writer.get_shape_color(obj) == "grey"