From 158b80c20394ee48a67d48f2c30eb73136645c3d Mon Sep 17 00:00:00 2001
From: JoschD <26184899+JoschD@users.noreply.github.com>
Date: Wed, 27 Nov 2024 13:26:58 +0100
Subject: [PATCH] return types

---
 scripts/shift_calculations.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/scripts/shift_calculations.py b/scripts/shift_calculations.py
index 40cdabb1..0e049484 100644
--- a/scripts/shift_calculations.py
+++ b/scripts/shift_calculations.py
@@ -84,7 +84,7 @@ def same_day(d1: datetime, d2: datetime) -> bool:
     return (d1.year == d2.year) and (d1.month == d2.month) and (d1.day == d2.day)
 
 
-def is_holiday(date: datetime):
+def is_holiday(date: datetime) -> bool:
     """ True is date is on a known holiday. """
     return any(same_day(date, h) for h in CERN_HOLIDAYS)
 
@@ -138,11 +138,11 @@ def calculate_shift_parts(start_time: datetime, end_time: datetime) -> dict[str,
     return time_split 
 
 
-def time_delta_to_hours(time_delta: timedelta):
+def time_delta_to_hours(time_delta: timedelta) -> float:
     return time_delta.total_seconds() / 3600
 
 
-def time_delta_to_shifts(time_delta: timedelta):
+def time_delta_to_shifts(time_delta: timedelta) -> float:
     return time_delta_to_hours(time_delta) / SHIFT_LENGTH
 
 
@@ -365,6 +365,8 @@ def plot_all_machines_in_year(
     fig.tight_layout()
     if output_path:
         fig.savefig(output_path)
+
+    return fig