Refactor shift constraints in tests to use new constraint classes

- Updated test cases in `test_workers.py` to replace dictionary-based constraints with instances of `PreShiftConstraint`, `PostShiftConstraint`, and `MaxShiftsPerWeekConstraint`.
- Added a new test `test_worker_assign_as_block_preference2` to verify worker preferences for block assignments.
- Marked tests as slow where appropriate to improve test suite performance.

gen_cons.py

@@ -3,7 +3,7 @@ import datetime
 import os
 import sys
 import time
-from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WorkerRequirement, days
+from rota.shifts import NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, WorkerRequirement, days
 import typer
 import subprocess
 import pandas as pd
@@ -335,7 +335,7 @@ def main(
             days=days,
             balance_offset=2,
             assign_as_block=False,
-            constraint=[
+            constraints=[
                 #{
                 #    "name": "max_shifts_per_week",
                 #    "options": 3,
@@ -349,9 +349,11 @@ def main(
             length=12.5,
             days=days[5:],
             balance_offset=1,
-            constraint=[
+            constraints=[
                 #{"name": "pre", "options": 2},
-                {"name": "post", "options": 1},
+                
+                #{"name": "post", "options": 1},
+                PostShiftConstraint(days=1, start_date="2025-11-24"),
             ],
             display_char="a",
             #force_assign_with=["oncall"]
@@ -380,7 +382,7 @@ def main(
             length=12.5,
             days=days[:5],
             balance_offset=1,
-            constraint=[
+            constraints=[
                 #{"name": "pre", "options": 1},
                 #{"name": "post", "options": 1},
             ],

rota/shifts.py

@@ -4,7 +4,7 @@ import itertools
 from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type, Literal, Optional
 import time
 
-from pydantic import BaseModel, ConfigDict, Extra, constr, model_validator, StrictInt, Field
+from pydantic import BaseModel, ConfigDict, Extra, constr, field_validator, model_validator, StrictInt, Field
 
 from pathlib import Path
 
@@ -59,18 +59,74 @@ SHIFT_BOUNDS = {
     "weekend_count": (0, 60),
 }
 
-VALID_SHIFT_CONSTRAINTS = Literal[
+#VALID_SHIFT_CONSTRAINTS = Literal[
-    "night",
+#    "night",
-    "pre",
+#    "pre",
-    "post",
+#    "post",
-    "balance_across_groups",
+#    "balance_across_groups",
-    "limit_grade_number",
+#    "limit_grade_number",
-    "minimum_grade_number",
+#    "minimum_grade_number",
-    "max_shifts_per_week",
+#    "max_shifts_per_week",
-    "max_shifts_per_week_block",
+#    "max_shifts_per_week_block",
-    "require_remote_site_presence",
+#    "require_remote_site_presence",
-    "require_remote_site_presence_week",
+#    "require_remote_site_presence_week",
-]
+#]
+
+class BaseShiftConstraint(BaseModel):
+    weeks: Optional[List[int]] = None
+    start_date: Optional[datetime.date] = None
+    end_date: Optional[datetime.date] = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+
+    @field_validator("weeks")
+    def weeks_none_if_dates_set(cls, weeks, values):
+        data = values.data
+        if (data.get("start_date") is not None or data.get("end_date") is not None) and weeks is not None:
+            raise ValueError("If start_date or end_date is set, weeks must not be set")
+        return weeks
+
+
+#class ShiftConstraint(BaseShiftConstraint):
+#    name: str
+#    options: Any = None
+
+class NightConstraint(BaseShiftConstraint):
+    options: Any = None
+
+class PreShiftConstraint(BaseShiftConstraint):
+    days: Any = None
+
+class PostShiftConstraint(BaseShiftConstraint):
+    days: Any = None
+
+class BalanceAcrossGroupsConstraint(BaseShiftConstraint):
+    """"""
+
+
+class MinSummedGradeByShiftsPerDayConstraint(BaseShiftConstraint):
+    shifts: List[str] = Field(..., description="List of shift names to sum grades over")
+    min_grade_sum: int = Field(..., description="Minimum total grade required")
+    days: Optional[List[str]] = None
+
+class LimitGradeNumberConstraint(BaseShiftConstraint):
+    grade: int 
+    max_number: int
+
+class MinimumGradeNumberConstraint(BaseShiftConstraint):
+    grade: int
+    min_number: int
+
+class RequireRemoteSitePresenceConstraint(BaseShiftConstraint):
+    site: str
+    required_number: int
+
+class MaxShiftsPerWeekConstraint(BaseShiftConstraint):
+    max_shifts: int | None = None
+
+class MaxShiftsPerWeekBlockConstraint(BaseShiftConstraint):
+    week_block : int
+    max_shifts: int | None = None
 
 def _pydantic_to_dict(obj):
     if isinstance(obj, list):
@@ -80,17 +136,6 @@ def _pydantic_to_dict(obj):
     return obj
 
 
-class ShiftConstraint(BaseModel):
-    name: VALID_SHIFT_CONSTRAINTS
-    options: bool | int | Dict | tuple = False
-
-class MinSummedGradeByShiftsPerDayConstraint(BaseModel):
-    shifts: List[str] = Field(..., description="List of shift names to sum grades over")
-    min_grade_sum: int = Field(..., description="Minimum total grade required")
-    weeks: Optional[List[int]] = None  # If None, applies to all weeks
-    days: Optional[List[str]] = None   # If None, applies to all days
-
-
 class WorkerRequirement(BaseModel):
     """
     start date is inclusive, end date is not
@@ -153,7 +198,9 @@ class SingleShift(BaseModel):
     force_as_block_unless_nwd: bool = False
     hard_constrain_shift: bool = True
     bank_holidays_only: bool = False
-    constraint: list[ShiftConstraint] = []
+    #constraint: list[ShiftConstraint] = []
+    constraints: List[BaseModel] = Field(default_factory=list)
+
     start_date: datetime.date | None = None
     end_date: datetime.date | None = None
     pair_proxy: bool = False
@@ -166,6 +213,11 @@ class SingleShift(BaseModel):
     )
 
     def __init__(self, **data: Any):
+        if "constraint" in data:
+            raise ValueError(
+                "The 'constraint' field on SingleShift has been deprecated and removed. "
+                "Please use the 'constraints' field with Pydantic constraint models instead."
+            )
         super().__init__(**data)
         # self.site = sites
         # self.name = name
@@ -187,13 +239,13 @@ class SingleShift(BaseModel):
 
         # self.assign_as_block = assign_as_block
 
-        self.constraint_options = {}
+        #self.constraint_options = {}
-        self.constraints = []
+        #self.constraints = []
-        for c in self.constraint:
+        #for c in self.constraint:
-            self.constraints.append(c.name)
+        #    self.constraints.append(c.name)
 
-            if c.options:
+        #    if c.options:
-                self.constraint_options[c.name] = c.options
+        #        self.constraint_options[c.name] = c.options
 
             # constraint = c.constraint
         #            match c:
@@ -206,7 +258,34 @@ class SingleShift(BaseModel):
             self.days = [self.days]
 
     def __str__(self):
-        return f"name: {self.name}, site: {self.sites}, length: {self.length}, balance weighting: {self.balance_weighting}, balance offset: {self.balance_offset}, hard_constrain: {self.hard_constrain_shift}, workers required : {self.workers_required}, rota on nwds {self.rota_on_nwds}, assign as block: {self.assign_as_block}, force as block: {self.force_as_block}, constraints: {self.constraints}, constraint_options: {self.constraint_options}"  # , total shifts: {self.total_shifts}"
+        return (
+            f"name: {self.name}, "
+            f"sites: {self.sites}, "
+            f"length: {self.length}, "
+            f"balance_weighting: {self.balance_weighting}, "
+            f"balance_offset: {self.balance_offset}, "
+            f"hard_constrain_shift: {self.hard_constrain_shift}, "
+            f"workers_required: {self.workers_required}, "
+            f"rota_on_nwds: {self.rota_on_nwds}, "
+            f"assign_as_block: {self.assign_as_block}, "
+            f"force_as_block: {self.force_as_block}, "
+            f"constraints: {self.constraints}"
+        )
+
+    def has_constraint(self, constraint: Type) -> bool:
+        for c in self.constraints:
+            if isinstance(c, constraint):
+                return True
+        return False
+
+    def get_constraint(self, constraint: Type) -> BaseShiftConstraint | None:
+        for c in self.constraints:
+            if isinstance(c, constraint):
+                return c
+        return None
+
+    def get_constraints(self, constraint: Type) -> list[BaseShiftConstraint]:
+        return [c for c in self.constraints if isinstance(c, constraint)]
 
     def get_shift_summary(self):
         """
@@ -641,7 +720,7 @@ class RotaBuilder(object):
                     (week, shift.name, site)
                     for site in self.sites
                     for week in self.weeks
-                    for shift in self.get_shifts_with_constraint("night")
+                    for shift in self.get_shifts_with_constraint(NightConstraint)
                 ),
                 # within=Binary,
                 within=NonNegativeIntegers,
@@ -653,7 +732,7 @@ class RotaBuilder(object):
                     (week, shift.name, site)
                     for site in self.sites
                     for week in self.weeks
-                    for shift in self.get_shifts_with_constraint("night")
+                    for shift in self.get_shifts_with_constraint(NightConstraint)
                 ),
                 domain=NonNegativeIntegers,
                 initialize=0,
@@ -663,7 +742,7 @@ class RotaBuilder(object):
                     (week, shift.name, site)
                     for site in self.sites
                     for week in self.weeks
-                    for shift in self.get_shifts_with_constraint("night")
+                    for shift in self.get_shifts_with_constraint(NightConstraint)
                 ),
                 domain=NonNegativeIntegers,
                 initialize=0,
@@ -675,7 +754,7 @@ class RotaBuilder(object):
                 (week, shift.name, site)
                 for site in self.sites
                 for week in self.weeks
-                for shift in self.get_shifts_with_constraint("balance_across_groups")
+                for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint)
             ),
             # within=Binary,
             within=NonNegativeIntegers,
@@ -687,7 +766,7 @@ class RotaBuilder(object):
                 (week, shift.name, site)
                 for site in self.sites
                 for week in self.weeks
-                for shift in self.get_shifts_with_constraint("balance_across_groups")
+                for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint)
             ),
             domain=NonNegativeIntegers,
             initialize=0,
@@ -697,7 +776,7 @@ class RotaBuilder(object):
                 (week, shift.name, site)
                 for site in self.sites
                 for week in self.weeks
-                for shift in self.get_shifts_with_constraint("balance_across_groups")
+                for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint)
             ),
             domain=NonNegativeIntegers,
             initialize=0,
@@ -1161,12 +1240,10 @@ class RotaBuilder(object):
                 <= limit
             )
 
-        for shift in self.get_shifts_with_constraint("limit_grade_number"):
+        for shift in self.get_shifts_with_constraint(LimitGradeNumberConstraint):
-            if not shift.constraint_options["limit_grade_number"]:
+            for constraint in shift.get_constraints(LimitGradeNumberConstraint):
-                raise ValueError(
+                grade = constraint.grade
-                    "Constraint option must be defined for 'limit_grade_number'"
+                limit = constraint.max_number
-                )
-            for grade in shift.constraint_options["limit_grade_number"]:
                 # self.model.limit_grades_constraint = Constraint(
                 setattr(
                     self.model,
@@ -1176,7 +1253,7 @@ class RotaBuilder(object):
                         # [shift for shift in self.get_shifts_with_constraint("limit_grade_number")],
                         [shift.name],
                         [grade],
-                        [shift.constraint_options["limit_grade_number"][grade]],
+                        [limit],
                         rule=limitGradesRule,
                         # name=f"limit_grades_constraint_{shift.name}_{grade}"
                     ),
@@ -1194,12 +1271,10 @@ class RotaBuilder(object):
                 >= limit
             )
 
-        for shift in self.get_shifts_with_constraint("minimum_grade_number"):
+        for shift in self.get_shifts_with_constraint(MinimumGradeNumberConstraint):
-            if not shift.constraint_options["minimum_grade_number"]:
+            constraint = shift.get_constraints(MinimumGradeNumberConstraint)[0]
-                raise ValueError(
+            grade, min_required = constraint.grade, constraint.min_number
-                    "Constraint option must be defined for 'minimum_grade_number'"
+
-                )
-            grade, min_required = shift.constraint_options["minimum_grade_number"]
             # self.model.limit_grades_constraint = Constraint(
             setattr(
                 self.model,
@@ -1231,12 +1306,11 @@ class RotaBuilder(object):
             )
 
         for shift in self.get_shifts_with_constraint(
-            "require_remote_site_presence_week"
+            RequireRemoteSitePresenceConstraint
         ):
-            site, required_number = shift.constraint_options[
+            constraint = shift.get_constraints(RequireRemoteSitePresenceConstraint)[0]
-                "require_remote_site_presence_week"
+            site, required_number = constraint.site, constraint.required_number
-            ]
+
-            # self.model.require_presence_at_site_overnight_rule = Constraint(
             setattr(
                 self.model,
                 f"require_remote_site_presence_week_{shift.name}",
@@ -1249,39 +1323,39 @@ class RotaBuilder(object):
                 ),
             )
 
-        def presenceAtRemoteSite(
+        #def presenceAtRemoteSite(
-            model, day, week, shift, required_site, required_number
+        #    model, day, week, shift, required_site, required_number
-        ):
+        #):
-            required_site_workers = [
+        #    required_site_workers = [
-                w for w in self.workers if required_site == w.remote_site
+        #        w for w in self.workers if required_site == w.remote_site
-            ]
+        #    ]
 
-            for w in required_site_workers:
+        #    for w in required_site_workers:
-                if (w.id, week, day, shift) not in model.works:
+        #        if (w.id, week, day, shift) not in model.works:
-                    return Constraint.Skip
+        #            return Constraint.Skip
 
-            return (
+        #    return (
-                sum(model.works[w.id, week, day, shift] for w in required_site_workers)
+        #        sum(model.works[w.id, week, day, shift] for w in required_site_workers)
-                >= required_number
+        #        >= required_number
-            )
+        #    )
 
-        for shift in self.get_shifts_with_constraint("require_remote_site_presence"):
+        #for shift in self.get_shifts_with_constraint("require_remote_site_presence"):
-            site, required_number = shift.constraint_options[
+        #    site, required_number = shift.constraint_options[
-                "require_remote_site_presence"
+        #        "require_remote_site_presence"
-            ]
+        #    ]
-            # self.model.require_presence_at_site_overnight_rule = Constraint(
+        #    # self.model.require_presence_at_site_overnight_rule = Constraint(
-            setattr(
+        #    setattr(
-                self.model,
+        #        self.model,
-                f"require_remote_site_presence_{shift.name}",
+        #        f"require_remote_site_presence_{shift.name}",
-                Constraint(
+        #        Constraint(
-                    [day for day in self.days],
+        #            [day for day in self.days],
-                    [week for week in self.weeks],
+        #            [week for week in self.weeks],
-                    [shift.name],
+        #            [shift.name],
-                    [site],
+        #            [site],
-                    [required_number],
+        #            [required_number],
-                    rule=presenceAtRemoteSite,
+        #            rule=presenceAtRemoteSite,
-                ),
+        #        ),
-            )
+        #    )
 
         # # Count the number of workers from each site on each night shift
         # # As 1 or 0 is optimum we can simply subtract 1 from the number
@@ -1309,7 +1383,7 @@ class RotaBuilder(object):
         for site in track(
             self.sites, description="Generating site balance constraints..."
         ):
-            for shift in self.get_shifts_with_constraint("balance_across_groups"):
+            for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint):
                 if site in shift.sites:
                     block = shift.name
                     for week in self.weeks:
@@ -2056,7 +2130,7 @@ class RotaBuilder(object):
                     == sum(
                         self.model.works[worker.id, week, day, shift.name]
                         for week, day, shift in self.get_week_day_shift_combinations_for_constraint(
-                            "night"
+                            NightConstraint
                         )
                         # for week, day in self.get_week_day_combinations_for_shift(shift)
                         # for shift in self.get_shifts_with_constraint("night")
@@ -2065,7 +2139,7 @@ class RotaBuilder(object):
 
                 night_shift_target_number = sum(
                     worker.shift_target_number[shift.name]
-                    for shift in self.get_shifts_with_constraint("night")
+                    for shift in self.get_shifts_with_constraint(NightConstraint)
                 )
 
                 self.model.constraints.add(
@@ -2217,7 +2291,7 @@ class RotaBuilder(object):
                         ):
                             continue
 
-                        if self.get_shifts_with_constraint("night"):
+                        if self.get_shifts_with_constraint(NightConstraint):
                             # Prevent nights more than once every n weeks
                             self.model.constraints.add(
                                 1
@@ -2227,7 +2301,7 @@ class RotaBuilder(object):
                                     ]
                                     for week in week_blocks
                                     for shift in self.get_shifts_with_constraint(
-                                        "night"
+                                        NightConstraint
                                     )
                                 )
                             )
@@ -2247,19 +2321,18 @@ class RotaBuilder(object):
                     )
 
             for constraint_shift in self.get_shifts_with_constraints(
-                "max_shifts_per_week_block"
+                MaxShiftsPerWeekBlockConstraint
             ):
-                constraint_options = constraint_shift.constraint_options[
+                constraint = constraint_shift.get_constraint(MaxShiftsPerWeekBlockConstraint)
-                    "max_shifts_per_week_block"
-                ]
 
                 # If not supplied, default to the number of days per shift
-                shift_number = constraint_shift.get_shift_number()
+                if constraint.max_shifts is not None:
-                if "shift_number" in constraint_options:
+                    shift_number = constraint.max_shifts
-                    shift_number = constraint_options["shift_number"]
+                else:
+                    shift_number = constraint_shift.get_shift_number()
 
                 for week_blocks in self.get_week_block_iterator(
-                    constraint_options["weeks"]
+                    constraint.week_block
                 ):
                     # Prevent shifts more than once every n weeks
                     self.model.constraints.add(
@@ -2289,11 +2362,19 @@ class RotaBuilder(object):
 
 
                 for constraint_shift in self.get_shifts_with_constraints(
-                    "max_shifts_per_week"
+                    MaxShiftsPerWeekConstraint
                 ):
+
+                    # If not supplied, default to the number of days per shift
+                    constraint = constraint_shift.get_constraint(MaxShiftsPerWeekConstraint)
+                    if constraint.max_shifts is not None:
+                        shift_number = constraint.max_shifts
+                    else:
+                        shift_number = constraint_shift.get_shift_number()
+
                     try:
                         self.model.constraints.add(
-                            constraint_shift.constraint_options["max_shifts_per_week"]
+                            shift_number
                             >= sum(
                                 self.model.works[
                                     worker.id, w, day, constraint_shift.name
@@ -2814,9 +2895,12 @@ class RotaBuilder(object):
                 # shift spans 7 days you may get >7 allocations in a row
                 # as it only checks for a different shift allocation
                 for constraint_shift in self.get_shifts_with_constraints(
-                    "pre"  # , week=week
+                    PreShiftConstraint
                 ):
-                    for n in range(0, constraint_shift.constraint_options["pre"]):
+                    constraint = constraint_shift.get_constraint(PreShiftConstraint)
+                    if week not in constraint.weeks:
+                        continue
+                    for n in range(0, constraint.days):
                         if day in constraint_shift.days:
                             # Only apply if worker can work this shift
                             if worker.site not in constraint_shift.sites:
@@ -2845,9 +2929,12 @@ class RotaBuilder(object):
                             )
 
                 for constraint_shift in self.get_shifts_with_constraints(
-                    "post"  # , week=week
+                    PostShiftConstraint
                 ):
-                    for n in range(0, constraint_shift.constraint_options["post"]):
+                    constraint = constraint_shift.get_constraint(PostShiftConstraint)
+                    if week not in constraint.weeks:
+                        continue
+                    for n in range(0, constraint.days):
                         if day in constraint_shift.days:
                             # Only apply if worker can work this shift
                             if worker.site not in constraint_shift.sites:
@@ -2877,7 +2964,7 @@ class RotaBuilder(object):
 
                 # Night constraint means we won't assign a shift the day before
                 # an unavailability
-                for constraint_shift in self.get_shifts_with_constraint("night"):
+                for constraint_shift in self.get_shifts_with_constraint(NightConstraint):
                     if (
                         worker.id,
                         # pweek,
@@ -3081,13 +3168,13 @@ class RotaBuilder(object):
                     * 20
                     # self.model.night_per_site2[week, block, site]
                     for week in self.weeks
-                    for shift in self.get_shifts_with_constraint("night")
+                    for shift in self.get_shifts_with_constraint(NightConstraint)
                     for site in self.sites
                 )
             else:
                 nights_site_balancing = 0
 
-            if self.get_shifts_with_constraint("balance_across_groups"):
+            if self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint):
                 block_site_balancing = sum(
                     (
                         self.model.shift_per_site_t1[week, shift.name, site]
@@ -3096,9 +3183,7 @@ class RotaBuilder(object):
                     * 2000
                     # self.model.night_per_site2[week, block, site]
                     for week in self.weeks
-                    for shift in self.get_shifts_with_constraint(
+                    for shift in self.get_shifts_with_constraint(BalanceAcrossGroupsConstraint)
-                        "balance_across_groups"
-                    )
                     for site in self.sites
                 )
             else:
@@ -3326,7 +3411,6 @@ class RotaBuilder(object):
             if return_empty_if_week_day_not_found:
                 return set()
             raise WeekDayNotFound(f"Week {week}, Day {day} not found in week_day_shifts_dict")
-        #print(self.week_day_shifts_dict)
         return self.week_day_shifts_dict[(week, day)]
 
     def get_week_day_shift_names_by_week_day(self, week, day: DayStr) -> set:
@@ -3363,7 +3447,7 @@ class RotaBuilder(object):
     def get_date_by_week_day(self, week: WeekInt, day: DayStr) -> datetime.date:
         return self.week_day_date_map[(week, day)]
 
-    def get_weeks_by_date_range(self, start_date: datetime.date, end_date: datetime.date) -> list[WeekInt]:
+    def get_weeks_by_date_range(self, start_date: datetime.date | None, end_date: datetime.date | None) -> list[WeekInt]:
         """Get a list of weeks that fall within a date range.
 
         Args:
@@ -3375,6 +3459,11 @@ class RotaBuilder(object):
         """
         # TODO: think about error if not monday start date?
         weeks = []
+        if start_date is None:
+            start_date = self.start_date
+        if end_date is None:
+            end_date = self.rota_end_date
+
         for week in self.weeks:
             week_start = self.get_date_by_week_day(week, "Mon")
             week_end = self.get_date_by_week_day(week, "Sun")
@@ -3471,7 +3560,6 @@ class RotaBuilder(object):
             # Check worker requirements
             if not isinstance(s.workers_required, int):
                 # Validate the worker requirements
-                print(s.workers_required)
                 for worker_requirement in s.workers_required:
                     if not isinstance(worker_requirement, WorkerRequirement):
                         raise InvalidShift(
@@ -3503,6 +3591,10 @@ class RotaBuilder(object):
             for site in list(s.sites):
                 self.sites.add(site)
 
+            for constraint in s.constraints:
+                if constraint.weeks is None:
+                    constraint.weeks = self.get_weeks_by_date_range(constraint.start_date, constraint.end_date)
+
         self.shift_counts = defaultdict(int)
         self.shift_worker_counts = defaultdict(int)
         for week, day in self.weeks_days_product:
@@ -3544,11 +3636,11 @@ class RotaBuilder(object):
 
         self.max_pre = 1
         self.max_post = 1
-        for shift in self.get_shifts_with_constraint("pre"):
+        for shift in self.get_shifts_with_constraint(PreShiftConstraint):
-            self.max_pre = max(shift.constraint_options["pre"], self.max_pre)
+            self.max_pre = max(shift.get_constraint(PreShiftConstraint).days, self.max_pre)
 
-        for shift in self.get_shifts_with_constraint("post"):
+        for shift in self.get_shifts_with_constraint(PostShiftConstraint):
-            self.max_post = max(shift.constraint_options["post"], self.max_post)
+            self.max_post = max(shift.get_constraint(PostShiftConstraint).days, self.max_post)
 
         # --- Check allowed_multi_shift_sets validity ---
         all_shift_names = set(s.name for s in self.shifts)
@@ -3708,7 +3800,7 @@ class RotaBuilder(object):
         return shifts
 
     def get_shifts_with_constraint(self, constraint) -> List[SingleShift]:
-        return [shift for shift in self.shifts if constraint in shift.constraints]
+        return [shift for shift in self.shifts if shift.has_constraint(constraint)]
 
     def get_shifts_with_constraints(
         self, *constraints, week: int | None = None
@@ -3720,7 +3812,7 @@ class RotaBuilder(object):
                 [
                     shift.name
                     for shift in self.get_shifts(week)
-                    if constraint in shift.constraints
+                    if shift.has_constraint(constraint)
                 ]
             )
 
@@ -4283,9 +4375,9 @@ class RotaBuilder(object):
                 remote_site = ""
                 if shift_name:
                     shift = self.get_shift_by_name(shift_name)
-                    if "require_remote_site_presence_week" in shift.constraints:
+                    if shift.has_constraint(RequireRemoteSitePresenceConstraint):
-                        remote_site = f" data-shift-remote-site='{shift.constraints['require_remote_site_presence_week'].options[0]}'"
+                        remote_site = f" data-shift-remote-site='{shift.get_constraint(RequireRemoteSitePresenceConstraint).site}'"
-                    if "night" in shift.constraints:
+                    if shift.has_constraint(NightConstraint):
                         css_class = " ".join((css_class, "night-shift"))
 
                 shift_tds.append(
@@ -4586,7 +4678,6 @@ class RotaBuilder(object):
         shifts = {}
 
         for worker in self.workers:
-            print("*", worker.name)
             shifts[worker.name] = len(
                 [i for i in self.get_worker_shift_list(worker) if i != ""]
             )
@@ -4703,7 +4794,6 @@ class RotaBuilder(object):
                 if c > 0:
                     l.append(f"{shift} ({c})")
                     total_shifts = total_shifts + c
-                # print(worker.id, shift)
                 timetable[worker.get_details()][shift] = c
             t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}")
         return "\n".join(t)
@@ -4732,7 +4822,6 @@ class RotaBuilder(object):
                 if c > 0:
                     l.append(f"{shift} ({c})")
                     total_shifts = total_shifts + c
-                # print(worker.id, shift)
                 shift_obj = self.get_shift_by_name(shift)
                 timetable[worker.get_details()][shift_obj.get_display_char()] = c
             t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}")

test/test_balance_shift_sites.py

@@ -1,6 +1,6 @@
 import pytest
 import datetime
-from rota.shifts import RotaBuilder, SingleShift, days
+from rota.shifts import BalanceAcrossGroupsConstraint, RotaBuilder, SingleShift, days
 from rota.workers import Worker
 
 @pytest.fixture
@@ -31,7 +31,7 @@ def test_balance_blocks_across_2_groups(demo_rota_balance_sites):
             balance_offset=40,
             workers_required=2,
             force_as_block=True,
-            constraint=[{"name": "balance_across_groups"}],
+            constraints=[BalanceAcrossGroupsConstraint()],
         ),
     )
     Rota.constraint_options["balance_nights_across_sites"] = False
@@ -76,7 +76,7 @@ def test_balance_blocks_across_2_groups_unbalanced(demo_rota_balance_sites):
             balance_offset=40,
             workers_required=3,
             force_as_block=True,
-            constraint=[{"name": "balance_across_groups"}],
+            constraints=[BalanceAcrossGroupsConstraint()],
         ),
     )
     Rota.constraint_options["balance_nights_across_sites"] = False
@@ -116,7 +116,7 @@ def test_balance_blocks_across_3_groups(demo_rota_balance_sites):
             balance_offset=40,
             workers_required=3,
             force_as_block=True,
-            constraint=[{"name": "balance_across_groups"}],
+            constraints=[BalanceAcrossGroupsConstraint()],
         ),
     )
     Rota.constraint_options["balance_nights_across_sites"] = False

test/test_general_balancing.py

@@ -1,7 +1,7 @@
 import datetime
 import pytest
 from pytest import approx
-from rota.shifts import RotaBuilder, SingleShift, days
+from rota.shifts import PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, days
 from rota.workers import Worker
 
 def generate_basic_rota(weeks_to_rota=10):
@@ -61,7 +61,7 @@ def test_weighted_shift_balancing():
             balance_weighting=10,
             workers_required=1,
             force_as_block=False,
-            constraint=[{"name": "pre","options": "2"}, {"name": "post","options": "2"}],
+            constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)]
         ),
         SingleShift(
             sites=("group1", "group2"),

test/test_limit_constraints.py

@@ -1,6 +1,6 @@
 import pytest
 import datetime
-from rota.shifts import RotaBuilder, SingleShift, days
+from rota.shifts import MinimumGradeNumberConstraint, RequireRemoteSitePresenceConstraint, RotaBuilder, SingleShift, days, LimitGradeNumberConstraint
 from rota.workers import Worker
 
 @pytest.fixture
@@ -33,7 +33,7 @@ def test_constraint_limit_grades(limit_constraint_rota):
             balance_offset=60,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "limit_grade_number", "options": {2: 1}}],
+            constraints=[LimitGradeNumberConstraint(grade=2, max_number=1)]
         ),
     )
     Rota.build_and_solve(options={"ratio": 0.01})
@@ -60,7 +60,7 @@ def test_constraint_limit_grades2(limit_constraint_rota):
             balance_offset=40,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "limit_grade_number", "options": {3: 1}}],
+            constraints=[LimitGradeNumberConstraint(grade=3, max_number=1)],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
@@ -96,7 +96,7 @@ def test_constraint_limit_grades3(limit_constraint_rota):
             balance_offset=60,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "limit_grade_number", "options": {2: 3, 3: 1}}],
+            constraints=[LimitGradeNumberConstraint(grade=2, max_number=3), LimitGradeNumberConstraint(grade=3, max_number=1)],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
@@ -124,7 +124,7 @@ def test_constraint_limit_grades4(limit_constraint_rota):
             balance_offset=40,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "limit_grade_number", "options": {2: 4, 3: 0}}],
+            constraints=[LimitGradeNumberConstraint(grade=2, max_number=4), LimitGradeNumberConstraint(grade=3, max_number=0)],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
@@ -143,7 +143,7 @@ def test_constraint_minimum_grades(limit_constraint_rota):
             balance_offset=40,
             workers_required=1,
             force_as_block=True,
-            constraint=[{"name": "minimum_grade_number", "options": (3, 1)}],
+            constraints=[MinimumGradeNumberConstraint(grade=3, min_number=1)],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
@@ -170,7 +170,7 @@ def test_constraint_minimum_grades2(limit_constraint_rota):
             balance_offset=40,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "minimum_grade_number", "options": (3, 3)}],
+            constraints=[MinimumGradeNumberConstraint(grade=3, min_number=3)],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
@@ -197,7 +197,7 @@ def test_constraint_minimum_grades3(limit_constraint_rota):
             balance_offset=40,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "minimum_grade_number", "options": (3, 0)}],
+            constraints=[MinimumGradeNumberConstraint(grade=3, min_number=0)],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
@@ -217,7 +217,7 @@ def test_constraint_minimum_grades_no_valid_worker(limit_constraint_rota):
             balance_offset=40,
             workers_required=4,
             force_as_block=True,
-            constraint=[{"name": "minimum_grade_number", "options": (4, 1)}],
+            constraints=[MinimumGradeNumberConstraint(grade=4, min_number=1)],
         ),
     )
     Rota.constraint_options["balance_shifts_quadratic"] = True
@@ -236,11 +236,8 @@ def test_constraint_require_remote_site_presence_week(limit_constraint_rota):
             balance_offset=40,
             workers_required=2,
             force_as_block=True,
-            constraint=[
+            constraints=[
-                {
+                RequireRemoteSitePresenceConstraint(site="group1", required_number=2)
-                    "name": "require_remote_site_presence_week",
-                    "options": ("group1", 2),
-                }
             ],
         ),
     )
@@ -272,11 +269,8 @@ def test_constraint_require_remote_site_presence_week2(limit_constraint_rota):
             balance_offset=20,
             workers_required=3,
             force_as_block=True,
-            constraint=[
+            constraints=[
-                {
+                RequireRemoteSitePresenceConstraint(site="group1", required_number=2)
-                    "name": "require_remote_site_presence_week",
-                    "options": ("group1", 2),
-                }
             ],
         ),
     )
@@ -304,11 +298,8 @@ def test_constraint_require_remote_site_presence_week3(limit_constraint_rota):
             balance_offset=20,
             workers_required=3,
             force_as_block=True,
-            constraint=[
+            constraints=[
-                {
+                RequireRemoteSitePresenceConstraint(site="group2", required_number=1)
-                    "name": "require_remote_site_presence_week",
-                    "options": ("group2", 1),
-                }
             ],
         ),
     )

test/test_night_unavailable.py

@@ -1,6 +1,6 @@
 import pytest
 import datetime
-from rota.shifts import RotaBuilder, SingleShift, days
+from rota.shifts import NightConstraint, RotaBuilder, SingleShift, days
 from rota.workers import Worker
 
 @pytest.fixture
@@ -45,7 +45,7 @@ def test_basic_assignment(demo_rota_night_unavailable):
             name="night_weekend",
             length=12.5,
             days=days[5:],
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
         ),
     )
     Rota.build_and_solve(options={"ratio": 0.00})
@@ -66,7 +66,7 @@ def test_assign_night_prior_to_unavailablity(demo_rota_night_unavailable):
             length=12.5,
             days=days[5:],
             workers_required=2,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
         ),
     )
     Rota.build_and_solve(options={"ratio": 0.00})
@@ -101,7 +101,7 @@ def test_assign_prior_to_unavailablity_non_night(demo_rota_night_unavailable):
             length=12.5,
             days=days[5:],
             workers_required=2,
-            constraint=[],
+            constraints=[],
         ),
     )
     Rota.build_and_solve(options={"ratio": 0.00})
@@ -135,7 +135,7 @@ def test_assign_split(demo_rota_night_unavailable):
             length=12.5,
             days=days[5:],
             workers_required=3,
-            constraint=[],
+            constraints=[],
         ),
     )
     Rota.build_and_solve(options={"ratio": 0.00})

test/test_nights2.py

@@ -1,6 +1,6 @@
 import datetime
 import pytest
-from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days
+from rota.shifts import InvalidShift, NightConstraint, NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, days
 from rota.workers import Worker
 
 def generate_basic_rota(weeks_to_rota=10):
@@ -35,7 +35,7 @@ def test_nights():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}, {"name": "pre", "options": 2}],
+            constraints=[NightConstraint(), PreShiftConstraint(days=2)],
             workers_required=2,
         ),
         SingleShift(
@@ -60,7 +60,7 @@ def test_nights_pre3():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+            constraints=[NightConstraint(), PreShiftConstraint(days=3)],
             workers_required=2,
         ),
         SingleShift(
@@ -84,7 +84,7 @@ def test_nights_pre3_2():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+            constraints=[NightConstraint(), PreShiftConstraint(days=3)],
             workers_required=1,
         ),
         SingleShift(
@@ -108,7 +108,7 @@ def test_nights_fail():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}, {"name": "pre", "options": 3}],
+            constraints=[NightConstraint(), PreShiftConstraint(days=3)],
             workers_required=2,
         ),
         SingleShift(
@@ -132,7 +132,7 @@ def test_nights2():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=2,
         ),
         SingleShift(
@@ -158,10 +158,10 @@ def test_nights_pre_wrap_around():
             length=12.5,
             days=days[:2],
             force_as_block=True,
-            constraint=[
+            constraints=[
-                {"name": "night"},
+                NightConstraint(),
-                {"name": "pre", "options": 5},
+                PreShiftConstraint(days=5),
-                {"name": "post", "options": 5},
+                PostShiftConstraint(days=5),
             ],
             workers_required=3,
         ),
@@ -195,7 +195,7 @@ def test_nights_max_frequency():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=2,
         ),
     )
@@ -213,7 +213,7 @@ def test_nights_max_frequency_fail():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=2,
         ),
     )
@@ -230,7 +230,7 @@ def test_nights_max_frequency3():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=1,
         ),
         SingleShift(
@@ -239,7 +239,7 @@ def test_nights_max_frequency3():
             length=12.5,
             days=days[:5],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=1,
         ),
     )
@@ -257,7 +257,7 @@ def test_nights_max_frequency4():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=1,
         ),
         SingleShift(
@@ -266,7 +266,7 @@ def test_nights_max_frequency4():
             length=12.5,
             days=days[:5],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=1,
         ),
     )
@@ -291,7 +291,7 @@ def test_nights_max_frequency_exclusions():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=2,
         ),
     )
@@ -310,7 +310,7 @@ def test_nights_max_frequency_exclusions2():
             length=12.5,
             days=days[5:],
             force_as_block=True,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=2,
         ),
     )
@@ -346,7 +346,7 @@ def test_nights_max_frequency_exclusions3():
             days=days[5:],
             force_as_block=True,
             balance_offset=10,
-            constraint=[{"name": "night"}],
+            constraints=[NightConstraint()],
             workers_required=2,
         ),
     )

test/test_remote.py

@@ -1,6 +1,6 @@
 
 import pytest
-from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
+from rota.shifts import NoWorkers, RequireRemoteSitePresenceConstraint, RotaBuilder, SingleShift, days
 
 import datetime
 from rota.workers import Worker
@@ -34,7 +34,7 @@ class TestRemoteRotas:
                 workers_required=1,
                 balance_offset=100,
                 force_as_block=False,
-                constraint=[{"name":"require_remote_site_presence_week", "options":("group2", 2)}],
+                constraints=[RequireRemoteSitePresenceConstraint(site="group2", required_number=2)],
             ),
         )
 
@@ -70,7 +70,7 @@ class TestRemoteRotas:
                 workers_required=1,
                 balance_offset=10,
                 force_as_block=True,
-                constraint=[{"name":"require_remote_site_presence_week", "options":("group2", 1)}],
+                constraints=[RequireRemoteSitePresenceConstraint(site="group2", required_number=1)],
             ),
             SingleShift(
                 sites=("group1", "group2"), name="b", length= 12.5, days=days[3:],
@@ -130,7 +130,7 @@ class TestRemoteRotas:
                 workers_required=1,
                 balance_offset=10,
                 force_as_block=False,
-                constraint=[{"name":"require_remote_site_presence_week", "options":("group2", 1)}],
+                constraints=[RequireRemoteSitePresenceConstraint(site="group2", required_number=1)],
             ),
             SingleShift(
                 sites=("group1", "group2"), name="b", length= 12.5, days=days[3:],

test/test_rota_clear.py

@@ -1,5 +1,5 @@
 import pytest
-from rota.shifts import RotaBuilder, SingleShift, days
+from rota.shifts import PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, days
 from rota.workers import Worker
 import datetime
 
@@ -59,7 +59,7 @@ def test_pre(demo_rota_clear):
             length=12.5,
             days=days[:5],
             workers_required=2,
-            constraint=[{"name": "pre", "options": 1}],
+            constraints=[PreShiftConstraint(days=1)],
         ),
         SingleShift(
             sites=("group1", "group2"),
@@ -74,7 +74,7 @@ def test_pre(demo_rota_clear):
             length=12.5,
             days=days[3],
             workers_required=1,
-            constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 1}],
+            constraints=[PreShiftConstraint(days=1), PostShiftConstraint(days=1)],
         ),
     )
 
@@ -123,12 +123,9 @@ def test_pre_post_clear(demo_rota_clear):
             days=days[:5],
             workers_required=2,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=1),
-                    "name": "pre",
+                PostShiftConstraint(days=1),
-                    "options": 1,
-                },
-                {"name": "post", "options": 1},
             ],
         ),
         SingleShift(
@@ -138,12 +135,9 @@ def test_pre_post_clear(demo_rota_clear):
             days=days[5:],
             workers_required=2,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=1),
-                    "name": "pre",
+                PostShiftConstraint(days=1),
-                    "options": 1,
-                },
-                {"name": "post", "options": 1},
             ],
         ),
     )
@@ -178,12 +172,9 @@ def test_pre_post_clear2(demo_rota_clear):
             days=days[:5],
             workers_required=2,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=2),
-                    "name": "pre",
+                PostShiftConstraint(days=2),
-                    "options": 2,
-                },
-                {"name": "post", "options": 2},
             ],
         ),
         SingleShift(
@@ -193,12 +184,9 @@ def test_pre_post_clear2(demo_rota_clear):
             days=days[5:],
             workers_required=2,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=2),
-                    "name": "pre",
+                PostShiftConstraint(days=2),
-                    "options": 2,
-                },
-                {"name": "post", "options": 2},
             ],
         ),
     )
@@ -237,12 +225,9 @@ def test_pre_post_clear_multi_group(demo_rota_clear):
             days=days[:5],
             workers_required=2,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=2),
-                    "name": "pre",
+                PostShiftConstraint(days=2),
-                    "options": 2,
-                },
-                {"name": "post", "options": 2},
             ],
         ),
         SingleShift(
@@ -252,12 +237,9 @@ def test_pre_post_clear_multi_group(demo_rota_clear):
             days=days[5:],
             workers_required=1,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=2),
-                    "name": "pre",
+                PostShiftConstraint(days=2),
-                    "options": 2,
-                },
-                {"name": "post", "options": 2},
             ],
         ),
         SingleShift(
@@ -267,12 +249,9 @@ def test_pre_post_clear_multi_group(demo_rota_clear):
             days=days[5:],
             workers_required=1,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=2),
-                    "name": "pre",
+                PostShiftConstraint(days=2),
-                    "options": 2,
-                },
-                {"name": "post", "options": 2},
             ],
         ),
     )
@@ -318,12 +297,9 @@ def test_pre_post_clear_force_assign(demo_rota_clear):
             days=days[:5],
             workers_required=1,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=1),
-                    "name": "pre",
+                PostShiftConstraint(days=1),
-                    "options": 1,
-                },
-                {"name": "post", "options": 1},
             ],
         ),
         SingleShift(
@@ -394,12 +370,7 @@ def test_pre_post_clear_force_assign2(demo_rota_clear):
             days=days[:5],
             workers_required=1,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                #{
-                #    "name": "pre",
-                #    "options": 1,
-                #},
-                #{"name": "post", "options": 1},
             ],
         ),
         SingleShift(
@@ -417,12 +388,7 @@ def test_pre_post_clear_force_assign2(demo_rota_clear):
             days=days[5:],
             workers_required=1,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                #{
-                #    "name": "pre",
-                #    "options": 2,
-                #},
-                #{"name": "post", "options": 2},
             ],
         ),
         SingleShift(
@@ -432,12 +398,9 @@ def test_pre_post_clear_force_assign2(demo_rota_clear):
             days=days[5:],
             workers_required=1,
             assign_as_block=False,  # global setting is off
-            constraint=[
+            constraints=[
-                {
+                PreShiftConstraint(days=1),
-                    "name": "pre",
+                PostShiftConstraint(days=1),
-                    "options": 1,
-                },
-                {"name": "post", "options": 1},
             ],
         ),
     )

test/test_rota_nights.py

@@ -1,8 +1,9 @@
 import datetime
 import pytest
 from pytest import approx
-from rota.shifts import RotaBuilder, SingleShift, days
+from rota.shifts import PostShiftConstraint, RotaBuilder, SingleShift, days
 from rota.workers import Worker
+from web.rota.shifts import PreShiftConstraint
 
 @pytest.fixture
 def demo_rota_nights():
@@ -48,7 +49,10 @@ def demo_rota_nights():
             length=12.5,
             days=days[:4],
             workers_required=1,
-            constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 2}],
+            constraints=[
+                PreShiftConstraint(days=1),
+                PostShiftConstraint(days=2),
+            ],
             force_as_block=True
         ),
         SingleShift(
@@ -56,7 +60,10 @@ def demo_rota_nights():
             name="night_weekend",
             length=12.5,
             days=days[4:],
-            constraint=[{"name": "pre", "options": 1}, {"name": "post", "options": 2}],
+            constraints=[
+                PreShiftConstraint(days=1),
+                PostShiftConstraint(days=2),
+            ],
             force_as_block=True,
         ),
         SingleShift(

test/test_workers.py

@@ -1,9 +1,10 @@
 import pytest
-from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days
+from rota.shifts import InvalidShift, MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, RotaBuilder, SingleShift, WarningTermination, WorkerRequirement, days
 
 import datetime
 from rota.workers import NonWorkingDays, NotAvailableToWork, Worker, generate_not_available_to_works
 from rota.workers import WorkRequests
+from web.rota.shifts import PreShiftConstraint
 
 def generate_basic_rota(
     weeks_to_rota=10, workers=0, start_date=datetime.date(2022, 3, 7)
@@ -577,9 +578,9 @@ def test_worker_requirement_and_force_block():
             days=days[:4],
             workers_required=wr,
             force_as_block=True,
-            constraint=[
+            constraints=[
-            {"name": "pre", "options": 2},
+                PreShiftConstraint(days=2),
-            {"name": "post", "options": 3},
+                PostShiftConstraint(days=3),
             ]
         ),
         SingleShift(
@@ -589,9 +590,9 @@ def test_worker_requirement_and_force_block():
             days=days[4:],
             workers_required=wr,
             force_as_block=True,
-            constraint=[
+            constraints=[
-            {"name": "pre", "options": 2},
+                PreShiftConstraint(days=2),
-            {"name": "post", "options": 3},
+                PostShiftConstraint(days=3),
             ]
         ),
     )
@@ -606,6 +607,7 @@ def test_worker_requirement_and_force_block():
         assert shift_string.count("cccc") == shift_string.count("c") / 4
         assert shift_string.count("ddd") == shift_string.count("d") / 3
 
+@pytest.mark.slow
 def test_split_shift():
     Rota = generate_basic_rota(workers=11, weeks_to_rota=22)
 
@@ -620,10 +622,10 @@ def test_split_shift():
             length=12.5,
             days=days[:4],
             workers_required=1,
-            constraint=[
+            constraints=[
-            {"name": "pre", "options": 2},
+                PreShiftConstraint(days=2),
-            {"name": "post", "options": 3},
+                PostShiftConstraint(days=3),
-            {"name": "max_shifts_per_week", "options": 1},
+                MaxShiftsPerWeekConstraint(max_shifts=1),
             ]
         ),
         SingleShift(
@@ -633,9 +635,9 @@ def test_split_shift():
             days=(days[4], days[6]),
             workers_required=1,
             assign_as_block=True,
-            constraint=[
+            constraints=[
-            {"name": "pre", "options": 2},
+                PreShiftConstraint(days=2),
-            {"name": "post", "options": 3},
+                PostShiftConstraint(days=3),
             ]
         ),
         SingleShift(
@@ -644,9 +646,9 @@ def test_split_shift():
             length=12.5,
             days=days[5],
             workers_required=1,
-            constraint=[
+            constraints=[
-            {"name": "pre", "options": 2},
+                PreShiftConstraint(days=2),
-            {"name": "post", "options": 3},
+                PostShiftConstraint(days=3),
             ]
         ),
     )
@@ -685,7 +687,7 @@ def test_locums_basic():
             days=days[:4],
             workers_required=1,
             
-            constraint=[
+            constraints=[
             #{"name": "pre", "options": 1},
             #{"name": "post", "options": 1},
             #{"name": "max_shifts_per_week", "options": 1},
@@ -734,7 +736,7 @@ def test_locums_nwds():
             days=days[:4],
             workers_required=1,
             
-            constraint=[
+            constraints=[
             #{"name": "pre", "options": 1},
             #{"name": "post", "options": 1},
             #{"name": "max_shifts_per_week", "options": 1},
@@ -751,7 +753,6 @@ def test_locums_nwds():
     assert Rota.results.solver.status == "ok"
 
     for worker in Rota.get_workers():
-        print(  Rota.get_worker_shift_list(worker))
         assert Rota.get_worker_shift_list(worker).count("a") == 4
 
         if worker.name in ("worker03", "worker04"):
@@ -782,7 +783,7 @@ def test_locums_no_availablity():
             days=days[:4],
             workers_required=1,
             
-            constraint=[
+            constraints=[
             #{"name": "pre", "options": 1},
             #{"name": "post", "options": 1},
             #{"name": "max_shifts_per_week", "options": 1},
@@ -794,6 +795,7 @@ def test_locums_no_availablity():
     with pytest.raises(WarningTermination):
         Rota.build_and_solve()
 
+@pytest.mark.slow
 def test_locums():
     Rota = generate_basic_rota(workers=7, weeks_to_rota=10)
 
@@ -832,10 +834,10 @@ def test_locums():
             days=days[:4],
             workers_required=1,
             
-            constraint=[
+            constraints=[
-            {"name": "pre", "options": 2},
+                PreShiftConstraint(days=2),
-            {"name": "post", "options": 3},
+                PostShiftConstraint(days=3),
-            {"name": "max_shifts_per_week", "options": 1},
+                MaxShiftsPerWeekConstraint(max_shifts=1),
             ]
         ),
         SingleShift(
@@ -845,9 +847,9 @@ def test_locums():
             days=(days[4], days[6]),
             workers_required=1,
             assign_as_block=True,
-            constraint=[
+            constraints=[
-            {"name": "pre", "options": 2},
+                PreShiftConstraint(days=2),
-            {"name": "post", "options": 3},
+                PostShiftConstraint(days=3),
             ]
         ),
         SingleShift(
@@ -856,9 +858,9 @@ def test_locums():
             length=12.5,
             days=days[5],
             workers_required=1,
-            constraint=[
+            constraints=[
-            {"name": "pre", "options": 2},
+                PreShiftConstraint(days=2),
-            {"name": "post", "options": 3},
+                PostShiftConstraint(days=3),
             ]
         ),
     )
@@ -912,11 +914,37 @@ def test_worker_assign_as_block_preference():
     shift_list_4 = Rota.get_worker_shift_list_string(worker4)
     assert shift_list_4.count("a") == 12, "Worker 4 should have at least 12 'a' shifts, but may not be grouped together"
 
+
+def test_worker_assign_as_block_preference2():
+    """Test that a worker's assign_as_block preference is respected over the global setting."""
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=8)
+
+    # Worker 1 prefers to have shift 'a' assigned as a block, worker 2 does not
+    workers = Rota.get_workers()
+    worker1 = workers[0]
+    worker2 = workers[1]
+    worker3 = workers[2]
+    worker4 = workers[3]
+
+
+    # Add a shift 'a' that is NOT globally set as assign_as_block
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+
     # Modify worker1's preference to prevent block 'a' shifts
     worker1.assign_as_block_preferences = {"a":-1}
     Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_assign_as_block_preference", folder="tests")
     shift_list_1 = Rota.get_worker_shift_list_string(worker1)
-    assert shift_list_1.count("aaaa") == 0, "Worker 1 should have all 'a' shifts grouped together (block)"
+    assert shift_list_1.count("aaaa") == 0, "Worker 1 should have less 'a' shifts grouped together (block)"
 
     worker1.assign_as_block_preferences = {"a":0}