Add support for avoiding shifts on specific dates for workers

- Introduced AvoidShiftOnDates model to define constraints for workers.
- Enhanced Worker class to register avoid-shift rules with the rota builder.
- Implemented logic in RotaBuilder to enforce these constraints during shift assignment.
- Added tests to verify functionality for avoiding shifts on single dates and date ranges.

leave.py

@@ -259,3 +259,9 @@ def load_leave(Rota):
 
 
 # load_leave()
+
+def load_academy(Rota):
+    """
+    Some trainees spend month blocks in the academy. This function loads that information from a google sheet that published it as a csv file.
+    """
+    url = "https://docs.google.com/spreadsheets/d/e/2PACX-1vQtNrp9F1fjtAh8vFX-W_R3RlD9H-2P_vCLMoHrVKR0e24yabbsdrD65VjwxoHlV1Qnatmw1NeghQHG/pub?gid=1113837782&single=true&output=csv"

output/tests/timetable.css

@@ -1 +1 @@
-/home/ross/proc-rota/output/timetable.css
+/home/ross/proc_rota/output/timetable.css

output/tests/timetable.js

@@ -1 +1 @@
-/home/ross/proc-rota/output/timetable.js
+/home/ross/proc_rota/output/timetable.js

rota_generator/shifts.py

@@ -300,10 +300,18 @@ class MaxShiftsPerWeekBlockConstraint(BaseShiftConstraint):
     max_shifts: int | None = None
 
 def _pydantic_to_dict(obj):
+    # Recursively convert pydantic models, lists and dicts to JSON-serializable structures
+    if obj is None:
+        return None
     if isinstance(obj, list):
         return [_pydantic_to_dict(x) for x in obj]
+    if isinstance(obj, dict):
+        return {k: _pydantic_to_dict(v) for k, v in obj.items()}
+    # datetime/date -> ISO string
+    if isinstance(obj, (datetime.date, datetime.datetime)):
+        return obj.isoformat()
     if hasattr(obj, "model_dump"):
-        return obj.model_dump()
+        return _pydantic_to_dict(obj.model_dump())
     return obj
 
 
@@ -582,6 +590,9 @@ class RotaConstraintOptions(BaseModel):
     # These are lists of dicts with keys like {"grades": [...], "weeks": [...], "shifts": [...]}
     avoid_shifts_by_grades: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning grades to shifts")
     avoid_shifts_by_worker_names: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning specific workers to shifts")
+    # Rules to avoid assigning specific workers to particular shifts on specific dates.
+    # Each entry is a dict with keys: {"names": [...], "shifts": [...], "start_date": date|None, "end_date": date|None, "dates": [date,...]|None}
+    avoid_shifts_by_worker_dates: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning specific workers to shifts on particular dates or date ranges")
     allocate_locum_shifts: bool = Field(True, description="Allow allocation of locum shifts")
     distribute_locum_shifts: bool = Field(True, description="Distribute locum shifts among available workers")
     balance_locum_shifts: bool = Field(False, description="Balance locum shifts similarly to normal shifts")
@@ -2311,6 +2322,61 @@ class RotaBuilder(object):
                     )
                     pass
 
+            # Date-based per-worker per-shift avoidance rules
+            for constraint_dict in self.constraint_options_model.avoid_shifts_by_worker_dates:
+                # validate shift names
+                if invalid_shifts := set(constraint_dict["shifts"]).difference(set(self.get_shift_names())):
+                    message = f"avoid shifts by worker/dates constraint contains a non existent shift: {invalid_shifts}"
+                    if self.ignore_valid_shifts:
+                        self.add_warning("Non existent shift", message)
+                        continue
+                    else:
+                        raise InvalidShift(message)
+
+                if worker.name not in constraint_dict["names"]:
+                    continue
+
+                # Build set of allowed (week,day) pairs that fall within the date selector
+                selected_weeks_days = set()
+
+                # explicit dates take precedence
+                if constraint_dict.get("dates"):
+                    date_set = set(constraint_dict["dates"])
+                    for wk, d in self.get_week_day_combinations():
+                        date = self.get_date_by_week_day(wk, d)
+                        if date in date_set:
+                            selected_weeks_days.add((wk, d))
+                else:
+                    sd = constraint_dict.get("start_date")
+                    ed = constraint_dict.get("end_date")
+                    # If neither provided, apply to whole rota
+                    if sd is None and ed is None:
+                        for wk, d in self.get_week_day_combinations():
+                            selected_weeks_days.add((wk, d))
+                    else:
+                        if sd is None:
+                            sd = self.start_date
+                        if ed is None:
+                            ed = self.rota_end_date
+                        for wk, d in self.get_week_day_combinations():
+                            date = self.get_date_by_week_day(wk, d)
+                            if sd <= date <= ed:
+                                selected_weeks_days.add((wk, d))
+
+                if not selected_weeks_days:
+                    # nothing to constrain
+                    continue
+
+                # Add constraint that forbids the given shifts for this worker on the selected dates
+                self.model.constraints.add(
+                    0
+                    == sum(
+                        self.model.works[worker.id, wk, d, shift]
+                        for wk, d, shift in self.get_all_shiftname_combinations()
+                        if (wk, d) in selected_weeks_days and shift in constraint_dict["shifts"]
+                    )
+                )
+
             for constraint_dict in self.constraint_options_model.avoid_shifts_by_worker_names:
                 if invalid_shifts := set(constraint_dict["shifts"]).difference(
                     set(self.get_shift_names())
@@ -4092,6 +4158,44 @@ class RotaBuilder(object):
             {"names": names, "weeks": weeks, "shifts": shifts}
         )
 
+    def add_avoid_shifts_for_workers_on_dates(
+        self,
+        names: Iterable[str],
+        shifts: Iterable[str],
+        start_date: datetime.date | str | None = None,
+        end_date: datetime.date | str | None = None,
+        dates: Iterable[datetime.date] | None = None,
+    ):
+        """Add a rule preventing the given workers from being assigned the given shifts
+        on the provided date range or explicit dates.
+
+        - `names`: iterable of worker names to which the rule applies.
+        - `shifts`: iterable of shift names to avoid.
+        - `start_date`/`end_date`: inclusive date range (accepts date or string "%d/%m/%Y" or "%d/%m/%y").
+        - `dates`: optional explicit iterable of `datetime.date` objects.
+        """
+        def _coerce_date(d):
+            if d is None:
+                return None
+            if isinstance(d, datetime.date):
+                return d
+            if isinstance(d, str):
+                for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"):
+                    try:
+                        return datetime.datetime.strptime(d, fmt).date()
+                    except ValueError:
+                        continue
+            raise ValueError(f"Cannot parse date: {d}")
+
+        sd = _coerce_date(start_date)
+        ed = _coerce_date(end_date)
+
+        entry = {"names": list(names), "shifts": list(shifts), "start_date": sd, "end_date": ed, "dates": None}
+        if dates is not None:
+            entry["dates"] = [d if isinstance(d, datetime.date) else _coerce_date(d) for d in dates]
+
+        self.constraint_options_model.avoid_shifts_by_worker_dates.append(entry)
+
     def add_grade_constraint_by_week(
         self, grades: Iterable[int], weeks: Iterable[int], shifts
     ):
@@ -4900,6 +5004,22 @@ class RotaBuilder(object):
                     timetable[worker.name][week][day] = assigned_shifts
         return timetable
 
+    def get_assigned_shifts_for_worker_on_day(self, worker_id: int, week: int, day: DayStr) -> list:
+        """Return a list of assigned SingleShift objects for the given worker on a specific week/day.
+
+        This is a small compatibility helper used by tests that expect objects rather than shift-name strings.
+        """
+        works = self.model.works
+        assigned = []
+        for shift_name in self.get_shift_names_by_week_day(week, day, return_empty_if_week_day_not_found=True):
+            try:
+                if works[worker_id, week, day, shift_name].value > 0.5:
+                    assigned.append(self.get_shift_by_name(shift_name))
+            except Exception:
+                # If variable missing or access error, skip
+                continue
+        return assigned
+
     def get_worker_timetable_brief(
         self, show_prefs=False, marker_every=30, show_unavailable=False
     ):

rota_generator/workers.py

@@ -111,6 +111,29 @@ class OutOfProgramme(BaseModel):
     reason: str = ""
 
 
+class AvoidShiftOnDates(BaseModel):
+    """Worker-level constraint: avoid listed shifts on given dates or date range."""
+    shifts: list[str]
+    start_date: Optional[datetime.date] = None
+    end_date: Optional[datetime.date] = None
+    dates: Optional[list[datetime.date]] = None
+
+    @field_validator("start_date", "end_date", mode="before")
+    def coerce_dates(cls, v):
+        # allow strings in common formats
+        if v is None:
+            return None
+        if isinstance(v, datetime.date):
+            return v
+        if isinstance(v, str):
+            for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"):
+                try:
+                    return datetime.datetime.strptime(v, fmt).date()
+                except Exception:
+                    continue
+        raise ValueError(f"Cannot parse date: {v}")
+
+
 class Worker(BaseModel):
     name: str
     site: str
@@ -154,6 +177,7 @@ class Worker(BaseModel):
     shift_fte_overrides: dict[str, int] = {} # Need checks to ensure shifts exist
 
     weekend_shift_target_number: int = 0
+    avoid_shifts_on_dates: list[AvoidShiftOnDates] = []
 
 
 
@@ -395,6 +419,22 @@ class Worker(BaseModel):
                 style="red",
             )
 
+        # Register any per-worker avoid-shift-on-dates rules with the rota builder
+        try:
+            for av in getattr(self, "avoid_shifts_on_dates", []):
+                entry = {
+                    "names": [self.name],
+                    "shifts": av.shifts,
+                    "start_date": av.start_date,
+                    "end_date": av.end_date,
+                    "dates": av.dates,
+                }
+                # append to rota-level list; the model builder will enforce these
+                if hasattr(Rota.constraint_options_model, "avoid_shifts_by_worker_dates"):
+                    Rota.constraint_options_model.avoid_shifts_by_worker_dates.append(entry)
+        except Exception:
+            Rota.add_warning("Worker/avoid_shifts_on_dates", f"Failed to register avoid_shifts_on_dates for {self.name}")
+
     def __lt__(self, other) -> bool:
         return (self.site, self.grade, self.fte_adj, self.name) < (
             other.site,

test/test_avoid_shifts_on_dates.py

@@ -0,0 +1,120 @@
+import datetime
+from rota_generator.shifts import RotaBuilder, SingleShift, days
+from rota_generator.workers import Worker, AvoidShiftOnDates
+
+
+def build_simple_rota():
+    start = datetime.date(2026, 4, 6)  # Monday
+    Rota = RotaBuilder(start, weeks_to_rota=4, name="test_rota")
+    Rota.add_shifts(
+        SingleShift(sites=("exeter",), name="night_weekday", length=12.25, days=days[:4], workers_required=1),
+        SingleShift(sites=("exeter",), name="plymouth_twilight", length=8, days=days[5:], workers_required=1),
+    )
+    Rota.build_shifts()
+    return Rota
+
+
+def test_avoid_single_date_for_worker():
+    Rota = build_simple_rota()
+    # create a worker who should avoid night_weekday on 2026-04-07
+    avoid = AvoidShiftOnDates(shifts=["night_weekday"], dates=[datetime.date(2026, 4, 7)])
+    w = Worker(name="Test Worker", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
+    Rota.add_worker(w)
+    Rota.build_workers()
+
+    # find week/day corresponding to 2026-04-07
+    weeks_days = [(wk, d) for wk, d in Rota.weeks_days_product if Rota.week_day_date_map[(wk, d)] == datetime.date(2026, 4, 7)]
+    assert len(weeks_days) == 1
+    wk, day = weeks_days[0]
+
+    # The model should mark works[worker.id, wk, day, 'night_weekday'] impossible -> check unavailable_to_work
+    assert (w.id, wk, day) not in Rota.unavailable_to_work
+
+    # Build and solve minimal model (but we can inspect constraints). Instead, check that the avoid rule was registered
+    found = False
+    for c in Rota.constraint_options_model.avoid_shifts_by_worker_dates:
+        if "names" in c and w.name in c["names"]:
+            found = True
+            # ensure the date is represented
+            assert c.get("dates") is not None and datetime.date(2026,4,7) in c.get("dates")
+    assert found
+
+
+def test_avoid_date_range_for_worker_applies_to_multiple_days():
+    Rota = build_simple_rota()
+    # Avoid night_weekday for 2026-04-06 through 2026-04-09
+    avoid = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="09/04/2026")
+    w = Worker(name="Range Worker", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
+    Rota.add_worker(w)
+    Rota.build_workers()
+
+    # Collect dates between start and end
+    sd = datetime.date(2026,4,6)
+    ed = datetime.date(2026,4,9)
+    dates_in_range = [d for d in Rota.get_date_range(sd, ed)]
+    assert len(dates_in_range) == 4
+
+    # Ensure the rota-level constraint was added
+    found = False
+    for c in Rota.constraint_options_model.avoid_shifts_by_worker_dates:
+        if w.name in c.get("names", []):
+            found = True
+            assert c.get("start_date") == sd
+            assert c.get("end_date") == ed
+    assert found
+
+
+def test_export_rota_for_visual_inspection():
+    Rota = build_simple_rota()
+    # add two workers
+    avoid = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="26/04/2026")
+    w1 = Worker(name="Export A", site="exeter", grade=3, fte=100, avoid_shifts_on_dates=[avoid])
+    w2 = Worker(name="Export B", site="exeter", grade=3, fte=100)
+
+
+
+    Rota.add_workers((w1, w2))
+
+    # Build model but do not solve, then export
+    Rota.name = "export_test_rota"
+    Rota.build_and_solve(solve=True)
+
+    assert Rota.results.solver.status == "ok", "Rota should be feasible"
+
+    # Next we check that worker w1 has no night_weekday shifts assigned between 06/04/2026 and 26/04/2026
+    sd = datetime.date(2026,4,6)
+    ed = datetime.date(2026,4,26)
+    for single_date in Rota.get_date_range(sd, ed):
+        wk_day = Rota.get_week_day_by_date(single_date)
+        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w1.id, wk_day[0], wk_day[1])
+        for sh in assigned_shifts:
+            assert sh.name != "night_weekday", f"Worker {w1.name} was assigned forbidden shift on {single_date}"
+
+    # We also check that w1 has 4 night shifts in total (the max possible in 4 weeks)
+    night_shift_count = 0
+    for wk, day in Rota.weeks_days_product:
+        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w1.id, wk, day)
+        for sh in assigned_shifts:
+            if sh.name == "night_weekday":
+                night_shift_count += 1
+    assert night_shift_count == 4, f"Worker {w1.name} should have 4 night shifts, has {night_shift_count}"
+
+    # And that w2 has 12 night shifts (no restrictions)
+    night_shift_count_w2 = 0
+    for wk, day in Rota.weeks_days_product:
+        assigned_shifts = Rota.get_assigned_shifts_for_worker_on_day(w2.id, wk, day)
+        for sh in assigned_shifts:
+            if sh.name == "night_weekday":
+                night_shift_count_w2 += 1
+    assert night_shift_count_w2 == 12, f"Worker {w2.name} should have 12 night shifts, has {night_shift_count_w2}"
+
+    # Next check that adding a short avoid to w2 results in an infeasible model
+    avoid2 = AvoidShiftOnDates(shifts=["night_weekday"], start_date="06/04/2026", end_date="07/04/2026")
+    w2.avoid_shifts_on_dates.append(avoid2)
+    Rota.build_and_solve(solve=True)
+
+    assert Rota.results.solver.termination_condition == "infeasible", "Rota should be infeasible due to avoid shifts on dates constraints"
+
+
+    Rota.export_rota_to_html("avoid_shifts_on_dates_test.html", folder="tests")
+