pyreverse: Add option for colored output

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
 

doc/whatsnew/2.10.rst

@@ -44,6 +44,8 @@ Extensions
 Other Changes
 =============
 
+* Pyreverse - add option to produce colored output.
+
 * Pyreverse - add output in PlantUML format
 
 * pylint does not crash with a traceback anymore when a file is problematic. It

pylint/pyreverse/dot_printer.py

@@ -35,6 +35,8 @@
 
 
 class DotPrinter(Printer):
+    DEFAULT_COLOR = "black"
+
     def __init__(
         self,
         title: str,
@@ -43,7 +45,6 @@ def __init__(
     ):
         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 @@ def emit_node(
         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 @@ def emit_node(
             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(

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 @@
             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",
         ),
     ),
     (
@@ -139,37 +139,56 @@
             help="create a *.<format> output file if format available.",
         ),
     ),
+    (
+        "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.",
+        ),
     ),
 )
 

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 @@ def __init__(self, config):
         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 @@ def set_printer(self, file_name: str, basename: str) -> None:
         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_color(obj) if self.config.colorized else "black",
         )
 
     def get_class_properties(self, obj: ClassEntity) -> NodeProperties:
@@ -123,10 +150,26 @@ def get_class_properties(self, obj: ClassEntity) -> NodeProperties:
             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_color(obj) if self.config.colorized else "black",
         )
         return properties
 
+    def get_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)

tests/pyreverse/conftest.py

@@ -23,6 +23,8 @@ def __init__(
         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 @@ def __init__(
         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
@@ -47,6 +51,14 @@ def default_config() -> PyreverseConfig:
     return PyreverseConfig()
 
 
+@pytest.fixture()
+def colorized_dot_config() -> PyreverseConfig:
+    return PyreverseConfig(
+        output_format="dot",
+        colorized=True,
+    )
+
+
 @pytest.fixture()
 def vcg_config() -> PyreverseConfig:
     return PyreverseConfig(
@@ -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:

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"];
+}

tests/pyreverse/data/classes_colorized.puml

@@ -0,0 +1,42 @@
+@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

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"];
+}

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