More refactoring and documenting

shifts.py

@@ -1,6 +1,10 @@
 import datetime
 import itertools
+from typing import List, Tuple
 
+ShiftName = str
+DayStr = str
+WeekInt = int
 
 days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
 
@@ -11,7 +15,9 @@ days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
 sites = ("truro", "exeter", "torquay", "barnstaple", "plymouth")
 
 
-class ShiftsClass(object):
+class SingleShift(object):
+    """Class to hold all details for a shift"""
+
     def __init__(
         self,
         sites,
@@ -32,8 +38,10 @@ class ShiftsClass(object):
 
 
 class ShiftCollection(object):
+    """Class to hold and manipulate shifts"""
+
     def __init__(self, weeks_to_rota=26):
-        self.shifts = []
+        self.shifts = []  # type List[SingleShift]
 
         self.weeks = [i for i in range(1, weeks_to_rota + 1)]
 
@@ -43,27 +51,56 @@ class ShiftCollection(object):
         self.rota_days_length = len(self.weeks) * 7
         self.rota_end_date = self.start_date + datetime.timedelta(self.rota_days_length)
 
-    def AddShift(self, shift):
-        """Add a shift to the collection"""
+    def add_shift(self, shift):
+        """Add a shift to the collection
+
+        :param SingleShift shift: Shift object
+
+        """
         self.shifts.append(shift)
         self.BuildShifts()
 
-    def AddShifts(self, *shifts):
+    def add_shifts(self, *shifts: SingleShift):
+        """Add multiple shifts
+        
+        Returns:
+            None
+        """
         self.shifts.extend(shifts)
         self.BuildShifts()
 
-    def GetShiftNamesByDay(self, day):
+    def GetShiftNamesByDay(self, day: DayStr):
+        """Returns the shifts required for a specific day
+
+        Returns:
+            set: Set of ShiftName
+
+        """
         return self.day_shifts_dict[day]
 
     def GetShiftByName(self, name):
+        """
+
+        :param name: 
+
+        """
         return self.shifts_by_name[name]
 
     def GetShiftLengthByName(self, name):
+        """Get the length of a shift
+        
+        Args:
+            name (str): Name of a shift
+        
+        Returns:
+            int: Length of the shift
+        """
         return self.shifts_by_name[name].length
 
     def BuildShifts(self):
+        """ """
         self.shifts_by_name = {}
-        self.shift_names = []
+        self.shift_names = []  # type: List[ShiftName]
 
         self.day_shifts_dict = {day: set() for day in days}
 
@@ -91,33 +128,85 @@ class ShiftCollection(object):
                     self.day_shiftclass_product.append((day, s))
 
     def GetAllDayShiftsAsNames(self):
+        """Returns a list of all required day / shift combinations
+        
+        Returns:
+            list: list of in the following format (str->day, str->shift)
+        """
         return self.day_shift_product
 
-    def GetAllDayShiftsAsClass(self):
+    def GetAllDayShiftsAsClass(self) -> List[Tuple[DayStr, SingleShift]]:
+        """Returns a list of all required day / shift combinations
+        
+        Returns:
+            list: list of in the following format (str->day, shiftObject->shift)
+        """
         return self.day_shiftclass_product
 
-    def GetAllShiftsAsNames(self):
+    def GetAllShiftsAsNames(self) -> List[Tuple[WeekInt, DayStr, ShiftName]]:
+        """Returns a list of all possible week / day / shift combinations
+        
+        Returns:
+            list: contains tuple of all possible week / day / shift combinations
+        """
         return self.week_day_shift_product
 
-    def GetAllShiftsAsClass(self):
+    def GetAllShiftsAsClass(self) -> List[SingleShift]:
+        """ """
         return self.week_day_shiftclass_product
 
-    def GetAllWeeksDays(self):
+    def GetAllWeeksDays(self) -> list:
+        """Returns a list of all week / day tuple combinations
+
+        Returns:
+            list: list of possible week day combinations
+                    [(1, "Mon"), (1, "Tue"), ... (n, "Fri")]
+        """
         return self.weeks_days_product
 
-    def GetShiftOptions(self):
+    def GetShiftOptions(self) -> List[SingleShift]:
+        """Returns a list of all the registered shifts
+        
+        Returns:
+            List[SingleShift]: list of registered shifts (as SingleShift)
+        """
         return self.shifts
 
-    def GetShiftNames(self):
+    def GetShiftNames(self) -> List[ShiftName]:
+        """Returns a list of all the registered shift names
+        
+        Returns:
+            List[ShiftName]: List of names of all available shifts
+        """ """ """
         return self.shift_names
 
     def GetWeeksInBlocks(self, block_length):
+        """Gets a two dimensional list of week blocks in specified length
+
+        e.g. block_length = 4
+            [
+                [ 1, 2, 3, 4],
+                [ 2, 3, 4, 5],
+                ...
+            ]
+        
+        Args:
+            block_length (int): length of blocks to create
+        
+        Returns:
+            list: two dimensional list containing weeks in blocks
+        """
         blocks = []
         for i in range(len(self.weeks)):
             blocks.append(self.weeks[i : i + block_length])
         return blocks
 
     def GetRequiredShiftsWorkersAndSites(self):
+        """Returns a list of all possible shifts, the workers required and site
+        
+        Returns:
+            list: list (week, day, shift name, workers required, shift site)
+        """
         l = []
         for week, day, shift in self.week_day_shiftclass_product:
             if day in shift.shift_days:
@@ -126,6 +215,13 @@ class ShiftCollection(object):
         return l
 
     def GetNotRequiredShifts(self):
+        """Returns a set of all possible shifts combinations that are not required
+
+        Includes those on days
+        
+        Returns:
+            list: (week, day, shift)
+        """
         l = set()
 
         for week in self.weeks:

workers.py

@@ -0,0 +1,60 @@
+import datetime
+
+from shifts import SingleShift, ShiftCollection, days, sites
+
+
+class Worker:
+    def __init__(self, Shifts, id, name, site, grade, fte=100, nwd=[], end_date=None, oop=None):
+        self.id = id
+        self.name = name
+        self.site = site
+        self.grade = grade
+        self.fte = fte
+        self.nwd = nwd
+        self.proportion_rota_to_work = 1
+
+        days_to_work = Shifts.rota_days_length
+
+        if end_date is None:
+            self.end_date = None
+        else:
+            self.end_date = datetime.datetime.strptime(end_date, '%Y/%m/%d').date()
+
+            if self.end_date > Shifts.rota_end_date:
+                self.end_date = Shifts.rota_end_date
+            else:
+                days_to_work = (self.end_date - Shifts.start_date).days
+
+                # add unavalabilities
+                for weeks_days in Shifts.weeks_days_product[days_to_work:]:
+                    week, day = weeks_days
+                    unavailable_to_work.add((self.id, week, day))
+        
+        if oop is not None:
+            start_oop, end_oop = oop
+            start_oop_date = datetime.datetime.strptime(start_oop, '%Y/%m/%d').date()
+            end_oop_date = datetime.datetime.strptime(end_oop, '%Y/%m/%d').date()
+
+            if end_oop_date > Shifts.rota_end_date:
+                end_oop_date = Shifts.rota_end_date
+            
+            oop_length = (end_oop_date - start_oop_date).days
+            days_to_work = days_to_work - oop_length
+
+            days_until_oop = (start_oop_date - Shifts.start_date).days
+
+            for weeks_days in Shifts.weeks_days_product[days_until_oop:days_until_oop+oop_length]:
+                week, day = weeks_days
+                unavailable_to_work.add((self.id, week, day))
+
+        
+        self.proportion_rota_to_work = days_to_work / Shifts.rota_days_length
+
+        # We had to adjust the full time equivalent for people who CCT / leave the rota early
+        self.fte_adj = self.fte * self.proportion_rota_to_work
+
+    def __lt__(self, other):
+        return (self.site, self.grade, self.name) < (other.site, other.grade, other.name)
+    
+    def get_details(self):
+        return "{} {} {}".format(self.id, self.site[0], self.grade)