From 9b1bbec189fa40b7cc6f2dfd808a04d719ad80ce Mon Sep 17 00:00:00 2001
From: Maren Philipps <55318391+m-philipps@users.noreply.github.com>
Date: Tue, 31 Oct 2023 14:45:26 +0100
Subject: [PATCH] Visualization: Fix `legend` argument checking for
 waterfall/parameter/history plots  (#1139)

- before the fix a `legend_error` would be thrown when the user supplied the legends as a tuple
- Now handles all data types
- raises `TypeError` for single and multiple results
- raises `ValueError` in case of multiple results when lengths of results and legends do not match


Co-authored-by: Dilan Pathirana <dilan.pathirana@uni-bonn.de>
Co-authored-by: Paul Jonas Jost <70631928+PaulJonasJost@users.noreply.github.com>
---
 pypesto/visualize/misc.py | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/pypesto/visualize/misc.py b/pypesto/visualize/misc.py
index d2fbbe348..739ae5e35 100644
--- a/pypesto/visualize/misc.py
+++ b/pypesto/visualize/misc.py
@@ -49,7 +49,7 @@ def process_result_list(
     """
     # check how many results were passed
     single_result = False
-    legend_error = False
+    legend_type_error = False
     if isinstance(results, list):
         if len(results) == 1:
             single_result = True
@@ -68,6 +68,10 @@ def process_result_list(
         # create list of legends for later handling
         if not isinstance(legends, list):
             legends = [legends]
+        try:
+            str(legends[0])
+        except TypeError:
+            legend_type_error = True
     else:
         # if more than one result is passed, we use one color per result
         colors = assign_colors_for_list(len(results), colors)
@@ -80,18 +84,19 @@ def process_result_list(
                 legends.append('Result ' + str(i_leg))
         else:
             # legends were passed by user: check length
-            if isinstance(legends, list):
+            try:
+                if isinstance(legends, str):
+                    legends = [legends]
                 if len(legends) != len(results):
-                    legend_error = True
-            else:
-                legend_error = True
-
-    # size of legend list and size of results does not match
-    if legend_error:
-        raise ValueError(
-            'List of results passed and list of labels do '
-            'not have the same length but should. Stopping.'
-        )
+                    raise ValueError(
+                        'List of results passed and list of labels do '
+                        'not have the same length.'
+                    )
+            except TypeError:
+                legend_type_error = True
+
+    if legend_type_error:
+        raise TypeError("Unexpected legend type.")
 
     return results, colors, legends