enh: return dict instead of list of gantt charts

rssched · Oct 4, 2024 · d4339b7 · d4339b7
1 parent f434ee5
commit d4339b7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/rssched/visualization/plot.py b/rssched/visualization/plot.py
@@ -6,7 +6,7 @@
 
 
 def generate_plots(response: Response, instance_name: str):
-    figures = plot_gantt_per_vehicle_type(response, instance_name)
+    figures = list(plot_gantt_per_vehicle_type(response, instance_name).values())
     figures.append(plot_active_events_over_time(response, instance_name))
     figures.append(plot_vehicle_utilization(response, instance_name))
     figures.extend(plot_fleet_efficiency(response, instance_name))

diff --git a/rssched/visualization/vehicle_type_gantt.py b/rssched/visualization/vehicle_type_gantt.py
@@ -4,8 +4,8 @@
 from rssched.visualization.colors import EVENT_TYPES
 
 
-def plot_gantt_per_vehicle_type(response: Response, instance_name: str):
-    figures = []
+def plot_gantt_per_vehicle_type(response: Response, instance_name: str) -> dict:
+    figures = {}
 
     for vehicle_type in response.schedule.fleet:
         vis_data = []
@@ -37,17 +37,15 @@ def plot_gantt_per_vehicle_type(response: Response, instance_name: str):
                         "Type": "DeadHeadTrip",
                     }
                 )
-        figures.append(
-            ff.create_gantt(
-                vis_data,
-                title=f"Rolling stock schedule: {vehicle_type.vehicle_type} (instance: {instance_name})",
-                colors=EVENT_TYPES,
-                index_col="Type",
-                show_colorbar=True,
-                group_tasks=True,
-                showgrid_x=True,
-                height=1200,
-            )
+        figures[vehicle_type.vehicle_type] = ff.create_gantt(
+            vis_data,
+            title=f"Rolling stock schedule: {vehicle_type.vehicle_type} (instance: {instance_name})",
+            colors=EVENT_TYPES,
+            index_col="Type",
+            show_colorbar=True,
+            group_tasks=True,
+            showgrid_x=True,
+            height=1200,
         )
 
     return figures