more updates

rota/shifts.py

@@ -23,6 +23,8 @@ from collections import defaultdict
 from io import StringIO
 import sys
 
+from rich.pretty import pprint
+
 ShiftName = str
 DayStr = str
 WeekInt = int
@@ -248,6 +250,8 @@ class RotaBuilder(object):
 
         self.results = None
 
+        self.ignore_valid_shifts = False
+
     def solve_model(
         self, solver: str = "cbc", use_neos: bool = False, options: dict = {}
     ):
@@ -279,6 +283,25 @@ class RotaBuilder(object):
         if not results.solver.status:
             sys.exit(0)
 
+    def solve_shifts_individually(self, options, folder="individual_shifts"):
+        shifts = self.shifts
+        
+        outcomes = {}
+
+        for shift in shifts:
+            self.ignore_valid_shifts = True
+            self.clear_shifts()
+            self.add_shift(shift)
+            self.build_and_solve(options)
+            self.export_rota_to_html(filename=shift.name, folder=folder)
+
+            outcomes[shift.name] = (self.results.solver.status, self.results.solver.termination_condition)
+            pprint(outcomes)
+
+
+        pprint(outcomes)
+
+
     def build_and_solve(
         self, options={"ratio": 0.1, "seconds": 1000, "threads": 10}, solve=True
     ):
@@ -640,7 +663,7 @@ class RotaBuilder(object):
         # Validate shifts are valid
         # work_request_sets = set([i[3] for i in self.work_requests])
         invalid_shifts = work_request_sets - set(self.get_shift_names())
-        if invalid_shifts:
+        if invalid_shifts and not self.ignore_valid_shifts:
             raise InvalidShift(
                 f"Invalid worker shift request [shift(s): ${invalid_shifts}]"
             )
@@ -1036,22 +1059,30 @@ class RotaBuilder(object):
 
             for week_blocks in self.get_week_block_iterator(4):
                 # Prevent more than n number shifts per 4 weeks
-                self.model.constraints.add(
-                    self.constraint_options["max_shifts_per_month"]
-                    >= sum(
-                        self.model.works[worker.id, week, day, shiftname]
-                        for week, day, shiftname in self.get_all_shiftname_combinations()
-                        if week in week_blocks
+                try:
+                    self.model.constraints.add(
+                        self.constraint_options["max_shifts_per_month"]
+                        >= sum(
+                            self.model.works[worker.id, week, day, shiftname]
+                            for week, day, shiftname in self.get_all_shiftname_combinations()
+                            if week in week_blocks
+                        )
                     )
-                )
+                except ValueError:
+                    # Occurs if there are no shifts within a defined block
+                    # TODO: test if this breaks (and we should check rathar than except)
+                    pass
 
             for constraint_dict in self.constraint_options["avoid_shifts_by_grades"]:
                 if invalid_shifts := set(constraint_dict["shifts"]).difference(
                     set(self.get_shift_names())
                 ):
-                    raise ValueError(
-                        f"avoid shifts by grade constraint contains a non existent shift: {invalid_shifts}"
-                    )
+                    if self.ignore_valid_shifts:
+                        continue # Skip if we don't want to raise an error
+                    else:
+                        raise InvalidShift(
+                            f"avoid shifts by grade constraint contains a non existent shift: {invalid_shifts}"
+                        )
                 if invalid_grades := set(constraint_dict["grades"]).difference(
                     set(self.get_worker_grades())
                 ):
@@ -1509,17 +1540,20 @@ class RotaBuilder(object):
                         )
                     )
 
-                self.model.constraints.add(
-                    self.constraint_options["max_shifts_per_week"]
-                    >= sum(
-                        self.model.works[worker.id, w, day, shiftname]
-                        # for shiftname in self.get_shift_names_by_week_day(week, day)
-                        # for day in self.days
-                        for w, day, shiftname in self.get_all_shiftname_combinations(
-                            week=week
+                try:
+                    self.model.constraints.add(
+                        self.constraint_options["max_shifts_per_week"]
+                        >= sum(
+                            self.model.works[worker.id, w, day, shiftname]
+                            # for shiftname in self.get_shift_names_by_week_day(week, day)
+                            # for day in self.days
+                            for w, day, shiftname in self.get_all_shiftname_combinations(
+                                week=week
+                            )
                         )
                     )
-                )
+                except ValueError:
+                    pass
 
                 # TODO: consider excluding shifts that span relevant period
                 if (
@@ -2350,6 +2384,9 @@ class RotaBuilder(object):
         """
         self.shifts.extend(shifts)
 
+    def clear_shifts(self) -> None:
+        self.shifts = []
+
     def get_shift_names_by_week_day(self, week, day: DayStr) -> set:
         """Returns the shifts required for a specific day
 
@@ -2701,8 +2738,14 @@ class RotaBuilder(object):
         return self.start_date + datetime.timedelta(weeks=week)
 
     # RESULTS
-    def export_rota_to_html(self, filename: str = "rota"):
-        output_file = Path("output", f"{filename}.html")
+    def export_rota_to_html(self, filename: str = "rota", folder = None):
+        if folder is not None:
+            output_file = Path("output", folder, f"{filename}.html")
+        else:
+            output_file = Path("output", f"{filename}.html")
+
+        # Make sure path exits
+        output_file.parent.mkdir(parents=True, exist_ok=True)
         with open(output_file, "w") as f:
             f.write(self.get_worker_timetable_html(True))