Seperate classes from notebook (helps with vscode intellisence)

shifts.py

@@ -0,0 +1,136 @@
+import datetime
+import itertools
+
+
+days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
+
+# weeks = [i for i in range(1,weeks_to_rota+1)]
+
+# weeks_days_product = list(itertools.product(weeks, days))
+
+sites = ("truro", "exeter", "torquay", "barnstaple", "plymouth")
+
+
+class ShiftsClass(object):
+    def __init__(
+        self,
+        sites,
+        name,
+        length,
+        shift_days,
+        balance_by_site=True,
+        workers_required=1,
+        rota_on_nwds=False,
+    ):
+        self.site = sites
+        self.name = name
+        self.length = length
+        self.shift_days = shift_days
+        self.balance_by_site = balance_by_site
+        self.workers_required = workers_required
+        self.rota_on_nwds = rota_on_nwds
+
+
+class ShiftCollection(object):
+    def __init__(self, weeks_to_rota=26):
+        self.shifts = []
+
+        self.weeks = [i for i in range(1, weeks_to_rota + 1)]
+
+        self.weeks_days_product = list(itertools.product(self.weeks, days))
+
+        self.start_date = datetime.date(2020, 9, 7)
+        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"""
+        self.shifts.append(shift)
+        self.BuildShifts()
+
+    def AddShifts(self, *shifts):
+        self.shifts.extend(shifts)
+        self.BuildShifts()
+
+    def GetShiftNamesByDay(self, day):
+        return self.day_shifts_dict[day]
+
+    def GetShiftByName(self, name):
+        return self.shifts_by_name[name]
+
+    def GetShiftLengthByName(self, name):
+        return self.shifts_by_name[name].length
+
+    def BuildShifts(self):
+        self.shifts_by_name = {}
+        self.shift_names = []
+
+        self.day_shifts_dict = {day: set() for day in days}
+
+        for s in self.shifts:
+            self.shifts_by_name[s.name] = s
+            self.shift_names.append(s.name)
+
+            for day in s.shift_days:
+                self.day_shifts_dict[day].add(s.name)
+
+        self.week_day_shift_product = []
+        self.week_day_shiftclass_product = []
+        for week, day in self.weeks_days_product:
+            for s in self.shifts:
+                if day in s.shift_days:
+                    self.week_day_shift_product.append((week, day, s.name))
+                    self.week_day_shiftclass_product.append((week, day, s))
+
+        self.day_shift_product = []
+        self.day_shiftclass_product = []
+        for s in self.shifts:
+            for day in days:
+                if day in s.shift_days:
+                    self.day_shift_product.append((day, s.name))
+                    self.day_shiftclass_product.append((day, s))
+
+    def GetAllDayShiftsAsNames(self):
+        return self.day_shift_product
+
+    def GetAllDayShiftsAsClass(self):
+        return self.day_shiftclass_product
+
+    def GetAllShiftsAsNames(self):
+        return self.week_day_shift_product
+
+    def GetAllShiftsAsClass(self):
+        return self.week_day_shiftclass_product
+
+    def GetAllWeeksDays(self):
+        return self.weeks_days_product
+
+    def GetShiftOptions(self):
+        return self.shifts
+
+    def GetShiftNames(self):
+        return self.shift_names
+
+    def GetWeeksInBlocks(self, block_length):
+        blocks = []
+        for i in range(len(self.weeks)):
+            blocks.append(self.weeks[i : i + block_length])
+        return blocks
+
+    def GetRequiredShiftsWorkersAndSites(self):
+        l = []
+        for week, day, shift in self.week_day_shiftclass_product:
+            if day in shift.shift_days:
+                l.append((week, day, shift.name, shift.workers_required, shift.site))
+
+        return l
+
+    def GetNotRequiredShifts(self):
+        l = set()
+
+        for week in self.weeks:
+            for day in days:
+                for shift in self.shifts:
+                    l.add((week, day, shift))
+
+        return l - set(self.GetAllShiftsAsClass())