Refactor gen.py and gen_proc.py to improve shift handling and add constraints; update leave.py to fetch new leave data source.

gen_cons.py

@@ -4,7 +4,7 @@ import datetime
 import os
 import sys
 import time
-from rota_generator.shifts import MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, WorkerRequirement, days, RotaConstraintOptions
+from rota_generator.shifts import MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, WorkerRequirement, days, RotaConstraintOptions, MaxShiftsPerWeekBlockConstraint
 import typer
 import subprocess
 import pandas as pd
@@ -404,6 +404,9 @@ def load_workers():
             previous_shifts=priors_map.get(initial, {}),
         )
 
+        if initial == "LB":
+            w.start_date = datetime.date(2026, 6, 15)
+
         #if initial == "L1":
         #    w.oop = [OutOfProgramme(
         #        start_date="2025-09-15",
@@ -482,7 +485,7 @@ def main(
         weeks_to_rota=weeks,
         balance_offset_modifier=bom,
         use_previous_shifts=True,
-        name="cons rota feb 2026 run 2",
+        name="cons rota feb 2026 run 5",
         allow_force_assignment_with_leave_conflict=True,
         constraint_options=RotaConstraintOptions(
             balance_weekends=True,
@@ -551,6 +554,7 @@ def main(
             constraints=[
                 PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"], exclude_days=("Sat", "Sun")),
                 MaxShiftsPerWeekConstraint(max_shifts=1),
+                MaxShiftsPerWeekBlockConstraint(week_block=3,max_shifts=1)
             ],
         ),
     )

gen_proc.py

@@ -2,7 +2,20 @@ import datetime
 import os
 import sys
 import time
-from rota_generator.shifts import NoWorkers, RotaBuilder, SingleShift, WorkerRequirement, days
+from rota_generator.shifts import (
+    NoWorkers,
+    RotaBuilder,
+    SingleShift,
+    WorkerRequirement,
+    days,
+    PreShiftConstraint,
+    PostShiftConstraint,
+    NightConstraint,
+    RequireRemoteSitePresenceConstraint,
+    LimitGradeNumberConstraint,
+    MinimumGradeNumberConstraint,
+    MaxShiftsPerWeekConstraint,
+)
 import typer
 import subprocess
 
@@ -35,8 +48,8 @@ def main(
     solve: bool = True,
     time_to_run: int = 60 * 60 * 4,
     ratio: float = 0.001,
-    start_date: datetime.datetime = "2025-09-01",
-    weeks: int = 16,
+    start_date: datetime.datetime = "2026-03-02",
+    weeks: int = 12,
     bom: int = 2,
 ):
     rota_start_date = start_date.date()
@@ -81,12 +94,7 @@ def main(
             length=12.5,
             days=days[:5],
             balance_offset=4,
-            constraint=[
-                {
-                    "name": "max_shifts_per_week",
-                    "options": 2,
-                }
-            ],
+            constraints=[MaxShiftsPerWeekConstraint(max_shifts=2)],
         ),
         SingleShift(
             sites=(
@@ -122,7 +130,7 @@ def main(
             days=days[:5],
             balance_offset=4,
             workers_required=1,
-            constraint=[{"name": "max_shifts_per_week", "options": 2}],
+            constraints=[MaxShiftsPerWeekConstraint(max_shifts=2)],
         ),
         SingleShift(
             sites=(
@@ -137,7 +145,7 @@ def main(
             balance_offset=3,
             rota_on_nwds=True,
             force_as_block=True,
-            constraint=[{"name": "post", "options": 2}, {"name": "pre", "options": 2}],
+            constraints=[PostShiftConstraint(days=2), PreShiftConstraint(days=2)],
             # constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2},],
         ),
         SingleShift(
@@ -155,7 +163,7 @@ def main(
             # rota_on_nwds=True,
             # force_as_block=True,
             assign_as_block=True,
-            constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
+            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
             # force_as_block_unless_nwd=True
         ),
         SingleShift(
@@ -166,7 +174,7 @@ def main(
             balance_offset=3,
             rota_on_nwds=True,
             force_as_block=True,
-            constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
+            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
             # force_as_block_unless_nwd=True
         ),
         SingleShift(
@@ -183,7 +191,7 @@ def main(
             workers_required=1,
             rota_on_nwds=True,
             force_as_block=True,
-            constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
+            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
         ),
         SingleShift(
             sites=(
@@ -198,7 +206,7 @@ def main(
             workers_required=1,
             rota_on_nwds=True,
             force_as_block=True,
-            constraint=[{"name": "pre", "options": 2}, {"name": "post", "options": 2}],
+            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
         ),
         SingleShift(
             sites=(
@@ -227,16 +235,13 @@ def main(
             workers_required=4,
             force_as_block=True,
             rota_on_nwds=True,
-            constraint=[
-                {"name": "night"},
-                {"name": "pre", "options": 2},
-                {"name": "post", "options": 2},
-                {
-                    "name": "require_remote_site_presence_week",
-                    "options": ("plymouth", 1),
-                },
-                {"name": "limit_grade_number", "options": {2: 1}},
-                {"name": "minimum_grade_number", "options": (4, 1)},
+            constraints=[
+                NightConstraint(),
+                PreShiftConstraint(days=2),
+                PostShiftConstraint(days=2),
+                RequireRemoteSitePresenceConstraint(site="plymouth", required_number=1),
+                LimitGradeNumberConstraint(grade=2, max_number=1),
+                MinimumGradeNumberConstraint(grade=4, min_number=1),
             ],
             #end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
         ),
@@ -257,16 +262,13 @@ def main(
             workers_required=4,
             force_as_block=True,
             rota_on_nwds=True,
-            constraint=[
-                {"name": "night"},
-                {"name": "pre", "options": 2},
-                {"name": "post", "options": 3},
-                {
-                    "name": "require_remote_site_presence_week",
-                    "options": ("plymouth", 1),
-                },
-                {"name": "limit_grade_number", "options": {2: 2}},
-                {"name": "minimum_grade_number", "options": (4, 1)},
+            constraints=[
+                NightConstraint(),
+                PreShiftConstraint(days=2),
+                PostShiftConstraint(days=3),
+                RequireRemoteSitePresenceConstraint(site="plymouth", required_number=1),
+                LimitGradeNumberConstraint(grade=2, max_number=2),
+                MinimumGradeNumberConstraint(grade=4, min_number=1),
             ],
             #end_date=datetime.datetime.strptime("2025-06-01", "%Y-%m-%d").date(),
         ),

leave.py

@@ -27,7 +27,7 @@ def load_leave(Rota):
         }
 
         if live_rota:
-            download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSfcQi5Qs__A8hE2CVvUA_6ULNJALrtZOBHPUZ3xyrxPif9obtQF2IqhioT_4nebjHV1Ac5iqhNtuq4/pub?gid=2024511103&single=true&output=csv", headers=headers)
+            download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vQtNrp9F1fjtAh8vFX-W_R3RlD9H-2P_vCLMoHrVKR0e24yabbsdrD65VjwxoHlV1Qnatmw1NeghQHG/pub?gid=1199196967&single=true&output=csv", headers=headers)
             #download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vSRx9VWXSlRubPyA0RhiI-Oqf5eHNYYEc6rFzlraDbR5_8qqr5g13-4uV-gn4u-TjZxiSMv1fBUaESq/pub?gid=814517272&single=true&output=csv", headers=headers)
             decoded_content = download.content.decode('utf-8')