From 81b2c85d61c8bdb986b054021ea5dd9a779e73fa Mon Sep 17 00:00:00 2001
From: Ross <xkjq@ymail.com>
Date: Wed, 17 Dec 2025 09:19:45 +0000
Subject: [PATCH] Add exclude_dates option to PreShiftConstraint and
 PostShiftConstraint

---
 gen_cons.py              |  4 ++--
 rota_generator/shifts.py | 30 ++++++++++++++++++++++++++----
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/gen_cons.py b/gen_cons.py
index de9b09c..4f4843d 100644
--- a/gen_cons.py
+++ b/gen_cons.py
@@ -341,7 +341,7 @@ def main(
         constraint_options=RotaConstraintOptions(
             balance_weekends=True,
             max_shifts_per_week=6,
-            max_days_per_week_block=[(4, 2), (6, 3)],
+            max_days_per_week_block=[(4, 2), (4, 3)],
         ),
     )
 
@@ -356,7 +356,7 @@ def main(
             balance_offset=2,
             assign_as_block=False,
             constraints=[
-                PreShiftConstraint(days=2, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun")),
+                PreShiftConstraint(days=2, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), exclude_dates=["2026-04-06", "2026-05-04", "2026-05-25"]),
                 PostShiftConstraint(days=1, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), allow_self=True),
             ],
         ),
diff --git a/rota_generator/shifts.py b/rota_generator/shifts.py
index a01f867..8b917fa 100644
--- a/rota_generator/shifts.py
+++ b/rota_generator/shifts.py
@@ -93,12 +93,14 @@ class PreShiftConstraint(BaseShiftConstraint):
     days: Optional[int] = None
     ignore_shifts: List[str] = []
     exclude_days: List[str] = []
+    exclude_dates: List[datetime.date] = []
     allow_self: bool = True
 
 class PostShiftConstraint(BaseShiftConstraint):
     days: Optional[int] = None
     ignore_shifts: List[str]  = []
     exclude_days: List[str] = []
+    exclude_dates: List[datetime.date] = []
     allow_self: bool = True
 
 class BalanceAcrossGroupsConstraint(BaseShiftConstraint):
@@ -3110,10 +3112,12 @@ class RotaBuilder(object):
                         ignore_shifts = constraint.ignore_shifts + [constraint_shift.name]
                     else:
                         ignore_shifts = constraint.ignore_shifts
-                    if week not in constraint.weeks:
+                    if constraint.weeks is not None and week not in constraint.weeks:
                         continue
+                    date_main = self.get_date_by_week_day(week, day)
                     for n in range(0, constraint.days):
-                        if day in constraint.exclude_days:
+                        # skip if main day excluded by weekday or specific date
+                        if day in getattr(constraint, "exclude_days", []) or date_main in getattr(constraint, "exclude_dates", []):
                             continue
                         if day in constraint_shift.days:
                             # Only apply if worker can work this shift
@@ -3125,6 +3129,14 @@ class RotaBuilder(object):
                                 ]
                             except KeyError:
                                 continue
+                            # compute the pre date and skip if it's in exclude_dates
+                            try:
+                                pre_date = self.get_date_by_week_day(pre_map[n][1], pre_map[n][2])
+                            except Exception:
+                                pre_date = None
+                            if pre_date is not None and pre_date in getattr(constraint, "exclude_dates", []):
+                                continue
+
                             self.model.constraints.add(
                                 1
                                 >= works
@@ -3150,10 +3162,12 @@ class RotaBuilder(object):
                         ignore_shifts = constraint.ignore_shifts + [constraint_shift.name]
                     else:
                         ignore_shifts = constraint.ignore_shifts
-                    if week not in constraint.weeks:
+                    if constraint.weeks is not None and week not in constraint.weeks:
                         continue
+                    date_main = self.get_date_by_week_day(week, day)
                     for n in range(0, constraint.days):
-                        if day in constraint.exclude_days:
+                        # skip if main day excluded by weekday or specific date
+                        if day in getattr(constraint, "exclude_days", []) or date_main in getattr(constraint, "exclude_dates", []):
                             continue
                         if day in constraint_shift.days:
                             # Only apply if worker can work this shift
@@ -3165,6 +3179,14 @@ class RotaBuilder(object):
                                 ]
                             except KeyError:
                                 continue
+                            # compute the post date and skip if it's in exclude_dates
+                            try:
+                                post_date = self.get_date_by_week_day(post_map[n][1], post_map[n][2])
+                            except Exception:
+                                post_date = None
+                            if post_date is not None and post_date in getattr(constraint, "exclude_dates", []):
+                                continue
+
                             self.model.constraints.add(
                                 1
                                 >= works