From 3a19a8ad652197f38bcbf6acae2378be10db3980 Mon Sep 17 00:00:00 2001
From: Galen Seilis <galen.seilis@seilis.ca>
Date: Fri, 15 Dec 2023 14:29:22 -0800
Subject: [PATCH] Docstrings and type hints for schedules.Schedule

Docstrings and type hints for schedules.Schedule
---
 ciw/schedules.py | 69 +++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 60 insertions(+), 9 deletions(-)

diff --git a/ciw/schedules.py b/ciw/schedules.py
index 918152f..044af70 100644
--- a/ciw/schedules.py
+++ b/ciw/schedules.py
@@ -1,11 +1,50 @@
+from typing import List, Tuple, Union, Generator, NoReturn
+
 class Schedule:
     """
-    A Schedule class, for storing information about server schedules and
-    generating the next shifts.
+    A Schedule class for storing information about server schedules and generating the next shifts.
+
+    Parameters
+    ----------
+    schedule : List[Tuple[int, float]]
+        A list of tuples representing shifts, where each tuple contains the number of servers and the shift date.
+    preemption : Union[bool, str], optional
+        Pre-emption option, should be either 'resume', 'restart', 'resample', or False. Default is False.
+
+    Attributes
+    ----------
+    schedule_type : str
+        Type of the schedule.
+    schedule_dates : List[float]
+        List of shift dates.
+    schedule_servers : List[int]
+        List of corresponding server numbers.
+    preemption : Union[bool, str]
+        Pre-emption option.
+    cyclelength : float
+        Length of the schedule cycle.
+
+    Methods
+    -------
+    initialise()
+        Initializes the generator object at the beginning of a simulation.
+    get_schedule_generator(boundaries, values)
+        A generator that yields the next time and number of servers according to a given schedule.
+    get_next_shift()
+        Updates the next shifts from the generator.
     """
-    def __init__(self, schedule, preemption=False):
+    def __init__(self, schedule: List[Tuple[int, float]], preemption: Union[bool, str] = False) -> NoReturn:
         """
-        Initialises the instance of the Schedule object
+        Initializes the instance of the Schedule object.
+
+        Parameters
+        ----------
+        schedule : List[Tuple[int, float]]
+            A list of tuples representing shifts, where each tuple contains
+            the number of servers and the shift date.
+        preemption : Union[bool, str], optional
+            Pre-emption option, should be either 'resume', 'restart', 'resample', or False.
+            Default is False.
         """
         if preemption not in [False, 'resume', 'restart', 'resample']:
             raise ValueError("Pre-emption options should be either 'resume', 'restart', 'resample', or False.")
@@ -15,17 +54,29 @@ def __init__(self, schedule, preemption=False):
         self.preemption = preemption
         self.cyclelength = self.schedule_dates[-1]
 
-    def initialise(self):
+    def initialise(self) -> NoReturn:
         """
-        Initialises the generator object at the beginning of a simulation
+        Initializes the generator object at the beginning of a simulation.
         """
         self.next_c = self.schedule_servers[0]
         self.schedule_generator = self.get_schedule_generator(self.schedule_dates, self.schedule_servers)
         self.get_next_shift()
 
-    def get_schedule_generator(self, boundaries, values):
+    def get_schedule_generator(self, boundaries: List[float], values:List[int]) -> Generator[Tuple[float, int], None, None]:
         """
         A generator that yields the next time and number of servers according to a given schedule.
+
+        Parameters
+        ----------
+        boundaries : List[float]
+            List of shift boundaries.
+        values : List[int]
+            List of corresponding server numbers.
+
+        Yields
+        ------
+        Tuple[float, int]
+            A tuple representing the next shift date and the number of servers.
         """
         num_boundaries = len(boundaries)
         index = 0
@@ -35,9 +86,9 @@ def get_schedule_generator(self, boundaries, values):
             index += 1
             yield date, values[index % num_boundaries]
 
-    def get_next_shift(self):
+    def get_next_shift(self) -> NoReturn:
         """
-        Updates the next shifts from the generator
+        Updates the next shifts from the generator.
         """
         self.c = self.next_c
         date, c = next(self.schedule_generator)