Add `--color-palette` option to `pyreverse` (#8223)

8 files changed, 96 insertions, 21 deletions

diff --git a/.pyenchant_pylint_custom_dict.txt b/.pyenchant_pylint_custom_dict.txt
index 47a222f8e..d3cce63e2 100644
--- a/.pyenchant_pylint_custom_dict.txt
+++ b/.pyenchant_pylint_custom_dict.txt
@@ -282,6 +282,7 @@ rstrip
 rtype
 runtime
 se
+seaborn
 sep
 setcomp
 shortstrings
diff --git a/doc/whatsnew/fragments/6738.feature b/doc/whatsnew/fragments/6738.feature
new file mode 100644
index 000000000..d6895df36
--- /dev/null
+++ b/doc/whatsnew/fragments/6738.feature
@@ -0,0 +1,3 @@
+`pyreverse` now supports custom color palettes with the `--color-palette` option.
+
+Closes #6738
diff --git a/pylint/pyreverse/main.py b/pylint/pyreverse/main.py
index 13669d5b4..5641aedb8 100644
--- a/pylint/pyreverse/main.py
+++ b/pylint/pyreverse/main.py
@@ -34,6 +34,26 @@ DIRECTLY_SUPPORTED_FORMATS = (
     "html",
 )
 
+DEFAULT_COLOR_PALETTE = (
+    "aliceblue",
+    "antiquewhite",
+    "aquamarine",
+    "burlywood",
+    "cadetblue",
+    "chartreuse",
+    "chocolate",
+    "coral",
+    "cornflowerblue",
+    "cyan",
+    "darkgoldenrod",
+    "darkseagreen",
+    "dodgerblue",
+    "forestgreen",
+    "gold",
+    "hotpink",
+    "mediumspringgreen",
+)
+
 OPTIONS: Options = (
     (
         "filter-mode",
@@ -174,6 +194,17 @@ OPTIONS: Options = (
         },
     ),
     (
+        "color-palette",
+        {
+            "dest": "color_palette",
+            "action": "store",
+            "default": DEFAULT_COLOR_PALETTE,
+            "metavar": "<color1,color2,...>",
+            "type": "csv",
+            "help": "Comma separated list of colors to use",
+        },
+    ),
+    (
         "ignore",
         {
             "type": "csv",
diff --git a/pylint/pyreverse/writer.py b/pylint/pyreverse/writer.py
index 68a49eea1..88e3d0535 100644
--- a/pylint/pyreverse/writer.py
+++ b/pylint/pyreverse/writer.py
@@ -34,27 +34,8 @@ class DiagramWriter:
         self.printer: Printer  # 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",
-            ]
-        )
+        # default colors are an adaption of the seaborn colorblind palette
+        self.available_colors = itertools.cycle(self.config.color_palette)
         self.used_colors: dict[str, str] = {}
 
     def write(self, diadefs: Iterable[ClassDiagram | PackageDiagram]) -> None:
diff --git a/pylint/testutils/pyreverse.py b/pylint/testutils/pyreverse.py
index 7a61ff5fe..20dce84dd 100644
--- a/pylint/testutils/pyreverse.py
+++ b/pylint/testutils/pyreverse.py
@@ -11,6 +11,8 @@ import sys
 from pathlib import Path
 from typing import NamedTuple
 
+from pylint.pyreverse.main import DEFAULT_COLOR_PALETTE
+
 if sys.version_info >= (3, 8):
     from typing import TypedDict
 else:
@@ -41,6 +43,7 @@ class PyreverseConfig(
         output_format: str = "dot",
         colorized: bool = False,
         max_color_depth: int = 2,
+        color_palette: tuple[str, ...] = DEFAULT_COLOR_PALETTE,
         ignore_list: tuple[str, ...] = tuple(),
         project: str = "",
         output_directory: str = "",
@@ -61,6 +64,7 @@ class PyreverseConfig(
         self.output_format = output_format
         self.colorized = colorized
         self.max_color_depth = max_color_depth
+        self.color_palette = color_palette
         self.ignore_list = ignore_list
         self.project = project
         self.output_directory = output_directory
diff --git a/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.puml b/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.puml
new file mode 100644
index 000000000..34dc814ff
--- /dev/null
+++ b/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.puml
@@ -0,0 +1,42 @@
+@startuml classes
+set namespaceSeparator none
+class "CheckerCollector" as custom_colors.CheckerCollector #red {
+  checker1
+  checker2
+  checker3
+}
+class "ElseifUsedChecker" as pylint.extensions.check_elif.ElseifUsedChecker #green {
+  msgs : dict
+  name : str
+  leave_module(_: nodes.Module) -> None
+  process_tokens(tokens: list[TokenInfo]) -> None
+  visit_if(node: nodes.If) -> None
+}
+class "ExceptionsChecker" as pylint.checkers.exceptions.ExceptionsChecker #yellow {
+  msgs : dict
+  name : str
+  options : tuple
+  open() -> None
+  visit_binop(node: nodes.BinOp) -> None
+  visit_compare(node: nodes.Compare) -> None
+  visit_raise(node: nodes.Raise) -> None
+  visit_tryexcept(node: nodes.TryExcept) -> None
+}
+class "StdlibChecker" as pylint.checkers.stdlib.StdlibChecker #yellow {
+  msgs : dict[str, MessageDefinitionTuple]
+  name : str
+  deprecated_arguments(method: str) -> tuple[tuple[int | None, str], ...]
+  deprecated_classes(module: str) -> Iterable[str]
+  deprecated_decorators() -> Iterable[str]
+  deprecated_methods() -> set[str]
+  visit_boolop(node: nodes.BoolOp) -> None
+  visit_call(node: nodes.Call) -> None
+  visit_functiondef(node: nodes.FunctionDef) -> None
+  visit_if(node: nodes.If) -> None
+  visit_ifexp(node: nodes.IfExp) -> None
+  visit_unaryop(node: nodes.UnaryOp) -> None
+}
+pylint.checkers.exceptions.ExceptionsChecker --* custom_colors.CheckerCollector : checker1
+pylint.checkers.stdlib.StdlibChecker --* custom_colors.CheckerCollector : checker3
+pylint.extensions.check_elif.ElseifUsedChecker --* custom_colors.CheckerCollector : checker2
+@enduml
diff --git a/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.py b/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.py
new file mode 100644
index 000000000..b46c97e75
--- /dev/null
+++ b/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.py
@@ -0,0 +1,10 @@
+from pylint.checkers.exceptions import ExceptionsChecker
+from pylint.checkers.stdlib import StdlibChecker
+from pylint.extensions.check_elif import ElseifUsedChecker
+
+
+class CheckerCollector:
+    def __init__(self):
+        self.checker1 = ExceptionsChecker(None)
+        self.checker2 = ElseifUsedChecker(None)
+        self.checker3 = StdlibChecker(None)
diff --git a/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.rc b/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.rc
new file mode 100644
index 000000000..ea30ad4a4
--- /dev/null
+++ b/tests/pyreverse/functional/class_diagrams/colorized_output/custom_colors.rc
@@ -0,0 +1,3 @@
+[testoptions]
+output_formats=puml
+command_line_args=-S --colorized --max-color-depth=2 --color-palette=red,green,yellow