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 <[email protected]>

Author: DudeNr33

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

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 <[email protected]>
 # Copyright (c) 2021 Marc Mueller <[email protected]>
 # Copyright (c) 2021 Mark Byrne <[email protected]>
+# 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
@@ -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/plantuml_printer.py

@@ -59,20 +59,25 @@ def emit_node(
             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,

pylint/pyreverse/printer.py

@@ -56,16 +56,25 @@ def __init__(
         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."""
 
     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(

pylint/pyreverse/vcg_printer.py

@@ -185,18 +185,9 @@
 
 
 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 @@ def _open_graph(self) -> None:
     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 @@ def emit_node(
             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 @@ def emit_edge(
     ) -> 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 @@ def _write_attributes(self, attributes_dict: Mapping[str, Any], **args) -> None:
                 ) 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]

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_shape_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_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)