Add allow_self option to PostShiftConstraint and PreShiftConstraint for flexible shift assignment

- 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.

.gitignore

@@ -13,3 +13,5 @@ logs/
 
 *.db
 web/rota
+
+cons/

.vscode/launch.json

@@ -20,5 +20,13 @@
             "program": "${workspaceFolder}/gen.py",
             "console": "integratedTerminal"
         },
+        {
+            "name": "Python: cons",
+            "type": "python",
+            "request": "launch",
+            "program": "${workspaceFolder}/gen_cons.py",
+            "console": "integratedTerminal"
+        },
+
     ]
 }

gen_cons.py

@@ -1,10 +1,12 @@
+from collections import defaultdict
 import datetime
 import os
 import sys
 import time
-from rota.shifts import NoWorkers, RotaBuilder, SingleShift, WorkerRequirement, days
+from rota.shifts import MaxShiftsPerWeekConstraint, NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift, WorkerRequirement, days
 import typer
 import subprocess
+import pandas as pd
 
 from rota.workers import (
     Worker,
@@ -15,12 +17,298 @@ from rota.workers import (
     OutOfProgramme,
 )
 
+from loguru import logger
 
 app = typer.Typer()
 
 
 sites = ("rota a", "rota b", "rota c")
 
+def extract_leave_and_rota_from_calender(calender_df):
+    """
+    Extract leave and rota assignments from calender_df.
+    - The first row contains worker names, starting from the 5th column (index 4).
+    - The second column (index 1) contains dates in dd/mm/yyyy format.
+    - The second column (index 1) also contains the rota/leave code for that date.
+    Returns a list of dicts: {"date": ..., "worker": ..., "rota": ...}
+    """
+    # Get worker names from the first row, starting from column 5 (index 4)
+    worker_names = [str(x).strip() for x in calender_df.columns[5:]]
+    # Iterate over each row (skip the header row)
+
+    # Get rota data from the second row (index 1), starting from column 5 (index 4)
+    # Extract the rota data for each worker from the first data row (index 0), columns 5 onwards
+    rota_row = calender_df.iloc[0]
+    rota_data = {}
+    for i, worker in enumerate(worker_names):
+        assignment = str(rota_row.iloc[5 + i]).strip() if len(rota_row) > 5 + i else ""
+        rota_data[worker] = f"rota {assignment.lower()}"
+
+    leave_requests = []
+    work_requests = []
+    # Process leave data from row 3 onwards (index 2 and up)
+    for idx, row in calender_df.iloc[3:].iterrows():
+        date_str = str(row.iloc[1]).strip().split(" ")[0]  # Get the date string from the second column (index 1)
+        if not date_str or date_str.lower() == "nan":
+            continue
+        try:
+            date = datetime.datetime.strptime(date_str, "%Y-%m-%d").date()
+        except Exception:
+            continue  # skip rows with invalid date
+        for i, worker in enumerate(worker_names):
+            if not worker or worker.lower() == "nan":
+                continue
+            leave_val = str(row.iloc[5 + i]).strip().lower() if len(row) > 5 + i else ""
+            if "no" in leave_val:
+                leave_requests.append({
+                    "date": date,
+                    "worker": worker,
+                })
+            elif "yes" in leave_val:
+                work_requests.append({
+                    "date": date,
+                    "worker": worker,
+                })
+            elif leave_val != "" and leave_val != "nan":
+                leave_requests.append({
+                    "date": date,
+                    "worker": worker,
+                })
+
+
+    #print(leave_requests)
+    return leave_requests, work_requests, rota_data
+
+def extract_shift_assignments_from_rota(rota_df):
+    """
+    Extract shift assignments from rota_df.
+    - Data starts from the 3rd row (index 2).
+    - The first column is the week start date.
+    - For each day (Mon-Sun), two columns: twilight initials, night initials.
+    Returns a list of dicts: {"week_start": ..., "day": ..., "shift": ..., "initial": ...}
+    """
+    days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
+    shifts = ["twilight", "oncall"]
+    assignments = []
+
+    for idx, row in rota_df.iloc[2:].iterrows():
+        week_start_str = str(row.iloc[0]).strip().split(" ")[0]  # Get the week start date from the first column (index 0)
+        try:
+            week_start = datetime.datetime.strptime(week_start_str, "%Y-%m-%d").date()
+        except Exception:
+            print(f"Invalid date format in row {idx}: {week_start_str}")
+            continue  # skip rows with invalid date
+
+        for i, day in enumerate(days):
+            for j, shift in enumerate(shifts):
+                if i < 5:
+                    # Mon-Fri: two columns per day (twilight, oncall)
+                    col_idx = 1 + i * 2 + j  # 1 for Mon twilight, 2 for Mon oncall, etc.
+                    weekend = False
+                else:
+                    # Sat/Sun: only one column for "weekend" shift (twilight), skip "night"
+                    weekend = True
+                    if j == 0:
+                        col_idx = 11 + (i - 5)  # 11 for Sat, 12 for Sun
+                    else:
+                        continue  # No night shift column for Sat/Sun
+                if col_idx >= len(row):
+                    continue
+                initial = str(row.iloc[col_idx]).strip()
+                if weekend:
+                    initial = initial.split("/")[0].strip()  # In case of "initial1/initial2", take the first one
+                if initial and initial.lower() != "nan" and initial.lower() != "nat":
+                    if weekend:
+                        shift = "weekend"
+                    assignments.append({
+                        "week_start": week_start,
+                        "week": idx-1,  # 1-indexed week number
+                        "day": day,
+                        "date": week_start + datetime.timedelta(days=i),
+                        "shift": shift,
+                        "initial": initial,
+                    })
+    # Build a mapping of assignments by worker initial
+    assignments_by_worker = {}
+    for assignment in assignments:
+        initial = assignment["initial"]
+        if initial not in assignments_by_worker:
+            assignments_by_worker[initial] = []
+        assignments_by_worker[initial].append(assignment)
+
+    return assignments, assignments_by_worker
+
+def load_workers():
+    # Path to the ODS file
+    ods_path = "cons/Oncall rota requests + New rota (first draft) - September 2025.ods"
+    
+    # Read all sheets
+    xls = pd.read_excel(ods_path, engine="odf", sheet_name=None)
+    
+    # Extract sheets
+    days_df = xls["Days"]
+    calender_df = xls["Calender"]
+    rota_df = xls["Rota"]
+
+
+    rota_b_path = "cons/Draft Rota B Sept 25-Feb 26.xlsx"
+    rota_b_df = pd.read_excel(rota_b_path, engine="openpyxl", sheet_name="Sheet1", header=None)
+    # Extract rota_b assignments: date in column A, worker in column B
+    rota_b_assignments = defaultdict(list)
+    for idx, row in rota_b_df.iterrows():
+        date_val = row.iloc[0]
+        worker_val = row.iloc[1] if len(row) > 1 else None
+        if pd.isna(date_val) or pd.isna(worker_val):
+            continue
+        try:
+            date = pd.to_datetime(date_val).date()
+        except Exception:
+            continue
+        worker = str(worker_val).strip()
+        if not worker or worker.lower() == "nan":
+            continue
+        rota_b_assignments[worker].append(date)
+    print(rota_b_assignments)
+    # You can now use rota_b_assignments as needed
+
+    # --- Example: Print the first few rows of each sheet ---
+    #print("Days sheet:")
+    #print(days_df.head())
+    weekday_cols = ["Mon", "Tue", "Wed", "Thu", "Fri"]
+    workers_days = []
+    for idx, row in days_df.iloc[7:].iterrows():
+        # Extract initials from the name column (assumed to be in brackets at the end)
+        text = str(row.iloc[1]).strip()
+        if "(" in text and ")" in text:
+            initial = text.split("(")[-1].split(")")[0].strip()
+        else:
+            raise ValueError(f"Invalid format in row {idx}: {text}")
+        name = text.split("(")[0].strip()
+        if not name or name.lower() == "nan":
+            continue
+        availability = {}
+        requests = {}
+        for i, day in enumerate(weekday_cols):
+            val = row.iloc[2 + i]
+            # You can adjust the logic below depending on your marking scheme (e.g. "Y", "Yes", "1", etc.)
+            available = not ("no" in str(val).strip().lower())
+            availability[day] = available
+        workers_days.append({"initial": initial, "name": name, "availability": availability, "requests": requests})
+
+
+    leave_requests, work_requests, rota_data = extract_leave_and_rota_from_calender(calender_df)
+
+    assignments, assignments_by_worker = extract_shift_assignments_from_rota(rota_df)
+    print(assignments)
+
+    workers = []
+    for worker_day in workers_days:
+        worker = worker_day["name"]
+        initial = worker_day["initial"]
+
+        start_date = "2025-09-01"  # Default start date, can be adjusted later
+
+        if initial == "ND":
+            start_date = "2025-09-29"  # Non-working day worker starts later
+        if initial == "MF":
+            start_date = "2026-01-15"  # Medical leave worker starts earlier
+        if initial == "AA":
+            start_date = "2025-11-17"  # Absent worker starts later
+
+        end_date = None
+        if initial == "L1":
+            end_date = "2025-11-17"  # L1 worker has a leave request
+        elif initial == "L2":
+            end_date = "2025-09-29"  # L2 worker has a leave request
+
+        non_working_days = []
+        for day, available in worker_day["availability"].items():
+            if not available:
+                non_working_days.append(NonWorkingDays(day=day, start_date=datetime.datetime(2025, 11, 17)))
+
+        w = Worker(
+            name=worker_day["initial"],
+            site=rota_data[worker],
+            grade=1,
+            start_date=start_date,
+            end_date=end_date,
+            # Add leave requests for this worker
+            not_available_to_work=[
+                NotAvailableToWork(
+                    date=req["date"],
+                    reason="leave",
+                ) for req in leave_requests if req["worker"] == worker
+            ],
+            work_requests=[
+                WorkRequests(
+                    date=req["date"],
+                    shift="*"
+                ) for req in work_requests if req["worker"] == worker
+            ],
+            nwds=non_working_days,
+        )
+
+        #if initial == "L1":
+        #    w.oop = [OutOfProgramme(
+        #        start_date="2025-09-15",
+        #        end_date="2025-10-14",
+        #        reason="even out shifts",
+        #    )]
+
+        if initial == "ND":
+            print(w.not_available_to_work)
+
+        if initial in assignments_by_worker:
+            # Add shift assignments for this worker
+            #if initial in ("RK", "HB", "DS", "L1", "L2", "L3", "GFM", "TS", "AB", "IR", "MF", "ND", "RG", "AA", "BR", "MS", "HA", "COD", "JL", "TB", "NHO"):
+                for assignment in assignments_by_worker[initial]:
+                    #if initial == "NHO":
+                    #    pass
+                    #    #if assignment["shift"] == "oncall":
+                    #    #    continue
+                    #    #if assignment["shift"] == "twilight":
+                    #    #    continue
+                    #    #1if assignment["shift"] == "weekend":
+                    #    #1    continue
+                    #print(assignment)
+                    w.force_assign_shift(
+                        week=assignment["week"],
+                        day=assignment["day"],
+                        shift_name=assignment["shift"],
+                    )
+        if initial in rota_b_assignments:
+                for date in rota_b_assignments.get(initial, []):
+                    w.force_assign_shift_by_date(
+                        date=date,
+                        shift_name="weekend b",
+                    )
+
+        if w.site == "rota a":
+
+            w.prefer_multi_shift_together = 10
+            w.allow_shifts_together("oncall", "twilight")
+            #w.add_force_assign_with("twilight", "oncall")
+            w.add_force_assign_with("weekend", "oncall")
+
+            if w.name in ("DS",):
+                w.add_hard_day_dependency({"Fri", "Sat", "Sun"}, start_date="2025-11-17")
+            else:
+                w.add_hard_day_dependency({"Fri", "Sun"}, start_date="2025-11-17")
+                w.add_hard_day_exclusion({"Sat", "Sun"}, start_date="2025-11-17")
+
+
+            #if w.name == "TB": 
+            #    w.add_hard_day_exclusion({"Sat", "Sun"}, start_date="2025-11-17")
+
+
+            #w.set_max_shifts
+                #w.set_max_shifts_per_week("weekend", 1)
+
+        workers.append(w)
+
+    return workers
+    
+
 
 @app.command()
 def main(
@@ -29,7 +317,7 @@ def main(
     time_to_run: int = 60 * 60,
     ratio: float = 0.001,
     start_date: datetime.datetime = "2025-09-01",
-    weeks: int = 22,
+    weeks: int = 24,
     bom: int = 1,
 ):
     rota_start_date = start_date.date()
@@ -39,13 +327,18 @@ def main(
         rota_start_date,
         weeks_to_rota=weeks,
         balance_offset_modifier=bom,
-        name="cons rota",
+        name="cons rota test",
+        allow_force_assignment_with_leave_conflict=True,
     )
     # Rota = RotaBuilder(start_date, weeks_to_rota=20, balance_offset_modifier=1)
     Rota.constraint_options["balance_weekends"] = True
-    Rota.constraint_options["max_weekend_frequency"] = 4
+    #Rota.constraint_options["max_weekend_frequency"] = 2
-    Rota.constraint_options["max_shifts_per_week"] = 20
+    Rota.constraint_options["max_shifts_per_week"] = 6
+    Rota.constraint_options["max_days_worked_per_week"] = 3
     # Rota.constraint_options["avoid_st2_first_month"] = True
+    Rota.constraint_options["max_days_per_week_block"] = [(4, 2), (6, 3)]  # (max_days, week_block)
+
+    #Rota.enable_unavailabilities = False
 
     Rota.add_shifts(
         SingleShift(
@@ -55,7 +348,9 @@ def main(
             days=days,
             balance_offset=2,
             assign_as_block=False,
-            constraint=[
+            constraints=[
+                PreShiftConstraint(days=2, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun")),
+                PostShiftConstraint(days=1, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), allow_self=True),
                 #{
                 #    "name": "max_shifts_per_week",
                 #    "options": 3,
@@ -68,41 +363,47 @@ def main(
             workers_required=1,
             length=12.5,
             days=days[5:],
-            balance_offset=2,
+            balance_offset=1,
-            constraint=[
+            constraints=[
                 #{"name": "pre", "options": 2},
+                
                 #{"name": "post", "options": 1},
+                PostShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"]),
             ],
             display_char="a",
             #force_assign_with=["oncall"]
         ),
         SingleShift(
-            sites=("rota b",),
+            sites=("rota b", "rota c"),
             name="weekend b",
             workers_required=1,
             length=12.5,
             days=days[5:],
             balance_offset=2,
-            constraint=[
+            #constraint=[
-                {"name": "pre", "options": 1},
+            #    {"name": "pre", "options": 1},
-                {"name": "post", "options": 1},
+            #    {"name": "post", "options": 1},
-            ],
+            #],
             display_char="b",
         ),
         SingleShift(
             sites=(
                 "rota a",
                 "rota b",
+                "rota d"
             ),
             name="twilight",
             workers_required=1,
             length=12.5,
             days=days[:5],
-            balance_offset=5,
+            balance_offset=1,
-            #constraint=[
+            constraints=[
-            #    {"name": "pre", "options": 1},
+                #PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"]),
-            #    {"name": "post", "options": 1},
+                PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"], exclude_days=("Sat", "Sun")),
-            #],
+                #PostShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall", "weekend a", "weekend b"], include_weekends=False),
+                MaxShiftsPerWeekConstraint(max_shifts=1),
+                #{"name": "post", "options": 1},
+            ],
         ),
     )
     #Rota.allow_shifts_together_for_all_workers(
@@ -120,74 +421,81 @@ def main(
     Rota.build_shifts()
     # Rota.add_grade_constraint_by_week([2], [1], Rota.get_shift_names())
 
-    for n in range(1, 12):
-        w = Worker(name=f"A{n:02d}", site="rota a", grade=1)
-        #w.prefer_multi_shift_together = 10
-        #w.allow_shifts_together("oncall", "twilight")
-        #w.allow_shifts_together("oncall", "weekend")
 
-        #w.add_hard_day_dependency("Fri", "Sat",  "Sun")
+    # Build workers
 
-        w.add_force_assign_with("twilight", "oncall")
+    workers = load_workers()
-        w.add_force_assign_with("weekend", "oncall")
+    Rota.add_workers(workers)
-        w.set_max_shifts_per_week("twilight", 1)
 
-        if n <= 2:
-            pass
-            w.add_hard_day_dependency("Fri", "Sat", "Sun")
-        else:
-            w.add_hard_day_dependency("Fri", "Sun")
-            #w.set_max_shifts_per_week("weekend", 1)
-        #    #w.add_hard_day_dependency("Sat", "Sun")
 
-        print(w)
+    #for n in range(1, 12):
+    #    w = Worker(name=f"A{n:02d}", site="rota a", grade=1)
+    #    #w.prefer_multi_shift_together = 10
+    #    #w.allow_shifts_together("oncall", "twilight")
+    #    #w.allow_shifts_together("oncall", "weekend")
 
-        Rota.add_worker(w)
+    #    #w.add_hard_day_dependency("Fri", "Sat",  "Sun")
 
-    w13 = Worker(name=f"A13", site="rota a", grade=1, start_date="2025-10-01")
+    #    w.add_force_assign_with("twilight", "oncall")
-    #w.prefer_multi_shift_together = 10
+    #    w.add_force_assign_with("weekend", "oncall")
-    #w.allow_shifts_together("oncall", "twilight")
+    #    w.set_max_shifts_per_week("twilight", 1)
-    #w.allow_shifts_together("oncall", "weekend")
 
-    #w.add_hard_day_dependency("Fri", "Sat",  "Sun")
+    #    if n <= 2:
+    #        pass
+    #        w.add_hard_day_dependency("Fri", "Sat", "Sun")
+    #    else:
+    #        w.add_hard_day_dependency("Fri", "Sun")
+    #        #w.set_max_shifts_per_week("weekend", 1)
+    #    #    #w.add_hard_day_dependency("Sat", "Sun")
 
-    w13.add_force_assign_with("twilight", "oncall")
+    #    print(w)
-    w13.add_force_assign_with("weekend", "oncall")
-    #w.set_max_shifts_per_week("twilight", 1)
 
-    Rota.add_worker(w13)
+    #    Rota.add_worker(w)
 
-    w13 = Worker(name=f"A14", site="rota a", grade=1, start_date="2025-12-01")
+    #w13 = Worker(name=f"A13", site="rota a", grade=1, start_date="2025-10-01")
-    #w.prefer_multi_shift_together = 10
+    ##w.prefer_multi_shift_together = 10
-    #w.allow_shifts_together("oncall", "twilight")
+    ##w.allow_shifts_together("oncall", "twilight")
-    #w.allow_shifts_together("oncall", "weekend")
+    ##w.allow_shifts_together("oncall", "weekend")
 
-    #w.add_hard_day_dependency("Fri", "Sat",  "Sun")
+    ##w.add_hard_day_dependency("Fri", "Sat",  "Sun")
 
-    w13.add_force_assign_with("twilight", "oncall")
+    #w13.add_force_assign_with("twilight", "oncall")
-    w13.add_force_assign_with("weekend", "oncall")
+    #w13.add_force_assign_with("weekend", "oncall")
-    #w.set_max_shifts_per_week("twilight", 1)
+    ##w.set_max_shifts_per_week("twilight", 1)
 
-    Rota.add_worker(w13)
+    #Rota.add_worker(w13)
 
-    for w in range(14, 24):
+    #w13 = Worker(name=f"A14", site="rota a", grade=1, start_date="2025-12-01")
-        w = Worker(
+    ##w.prefer_multi_shift_together = 10
-            name=f"B{w}",
+    ##w.allow_shifts_together("oncall", "twilight")
-            site="rota b",
+    ##w.allow_shifts_together("oncall", "weekend")
-            grade=1,
-            nwds=[
-                NonWorkingDays(
-                    day="Fri",
-                )
-            ],
-        )
-        #w.add_hard_day_exclusion("Sat", "Sun")
-        w.set_max_shifts_per_week("twilight", 2)
-        w.set_max_shifts_per_week("weekend b", 1)
 
-        print(w)
+    ##w.add_hard_day_dependency("Fri", "Sat",  "Sun")
 
-        Rota.add_worker(w)
+    #w13.add_force_assign_with("twilight", "oncall")
+    #w13.add_force_assign_with("weekend", "oncall")
+    ##w.set_max_shifts_per_week("twilight", 1)
+
+    #Rota.add_worker(w13)
+
+    #for w in range(14, 24):
+    #    w = Worker(
+    #        name=f"B{w}",
+    #        site="rota b",
+    #        grade=1,
+    #        nwds=[
+    #            NonWorkingDays(
+    #                day="Fri",
+    #            )
+    #        ],
+    #    )
+    #    #w.add_hard_day_exclusion("Sat", "Sun")
+    #    w.set_max_shifts_per_week("twilight", 2)
+    #    w.set_max_shifts_per_week("weekend b", 1)
+
+    #    print(w)
+
+    #    Rota.add_worker(w)
 
     # Rota.build_workers()
     # Rota.build_model()
@@ -196,7 +504,7 @@ def main(
     # solver_options = {"seconds": time_to_run, "threads": 10}
 
     # start_time = time.time()
-    Rota.build_and_solve(solver_options, export=True, solve=solve, solver="appsi_highs")
+    Rota.build_and_solve(solver_options, export=True, export_with_timestamp=False, solve=solve, solver="appsi_highs")
 
     # Rota.solve_shifts_by_block(solver_options, block_length=13)
     # Rota.solve_shifts_individually(solver_options)

output/timetable.css

@@ -31,14 +31,14 @@ table {
 
 .table-div {
     /* overflow-x: auto; */
-    width: 80%;
+    /* width: 80%; */
     overflow-x: scroll;
     margin-left: 200px;
     overflow-y: visible;
     padding: 0;
 }
 
-.table-div th {
+#main-table th {
     text-align: center;
     white-space: nowrap;
     transform: rotate(90deg) translateX(-100px);
@@ -47,12 +47,12 @@ table {
 
 }
 
-.table-div th p {
+#main-table th p {
     margin: 0 -100%;
     display: inline-block;
 }
 
-.table-div th p:before {
+#main-table th p:before {
     content: '';
     width: 0;
     padding-top: 110%;
@@ -61,21 +61,21 @@ table {
     vertical-align: middle;
 }
 
-.table-div td,
+#main-table td,
 th {
     max-width: 10px;
     padding: 0px;
 }
 
-.table-div tr {
+#main-table tr {
     padding: 0px;
 }
 
-.table-div tr:hover {
+#main-table tr:hover {
     background-color: lightblue;
 }
 
-.table-div tr:hover .unavailable {
+#main-table tr:hover .unavailable {
 
     background-color: lightgray;
 }
@@ -96,7 +96,7 @@ th {
 
 }
 
-th.bank-holiday {
+#main-table th.bank-holiday {
     color: darkblue;
 
 }
@@ -118,7 +118,8 @@ th.bank-holiday {
     color: #0074D9;
 }
 
-.torquay, .torbay {
+.torquay,
+.torbay {
     color: #5FA8E8;
 }
 
@@ -186,10 +187,8 @@ th.bank-holiday {
     color: lightcoral;
 }
 
-.th {}
+#main-table th.worker,
-
+#main-table td.worker {
-.table-div th.worker,
-.table-div td.worker {
     position: absolute;
     width: 200px;
     max-width: 200px;
@@ -209,7 +208,7 @@ th.bank-holiday {
     max-width: auto;
 }
 
-.header th {
+#main-table.header th {
     max-width: 200px;
 
 }
@@ -218,7 +217,7 @@ th.bank-holiday {
     border: 1px dotted orange;
 } */
 
-th.site-title {
+#main-table th.site-title {
     transform: unset;
 }
 
@@ -250,6 +249,7 @@ td.large-shift-diff-pos::before {
     content: "*";
     color: red;
 }
+
 td.large-shift-diff-neg::before {
     content: "*";
     color: green;
@@ -259,6 +259,7 @@ tr.large-shift-diff-pos td {
     border-bottom: solid 1px red;
     border-top: solid 1px red;
 }
+
 tr.large-shift-diff-neg td {
     border-bottom: solid 1px green;
     border-top: solid 1px green;
@@ -279,11 +280,9 @@ table.transposed td {
     border: 1px solid black;
     max-width: 100px;
 } */
-table.transposed {
+table.transposed {}
-}
 
-table.transposed tr {
+table.transposed tr {}
-}
 
 table.transposed th,
 table.transposed td {
@@ -326,3 +325,43 @@ table.transposed th.bank-holiday {
 .multi-shift {
     line-height: 1;
 }
+
+button.selected {
+    background: #0074d9;
+    color: #fff;
+}
+
+#main-table th.col-hover {
+    background: #ffe066 !important;
+    color: #222 !important;
+}
+
+.highlight-day-btn.highlighted {
+    background: #ffe066 !important;
+    color: #222 !important;
+    border: 2px solid #bfa600 !important;
+}
+
+.highlight-day-btn.highlighted {
+    background: #ffe066 !important;
+    color: #222 !important;
+    border: 2px solid #bfa600 !important;
+}
+
+#main-table td.day-highlighted,
+#main-table th.day-highlighted {
+    /* background is set inline for custom color */
+}
+
+#display-settings-modal {
+    transition: box-shadow 0.2s;
+}
+
+#main-table td.force-assigned.force-assigned-highlight {
+    background-color: #ffe066 !important;
+    border: 2px solid #ff8800 !important;
+}
+
+#linear-shift-timetable td, #linear-shift-timetable th {
+    max-width: unset;
+}

output/timetable.js

@@ -18,6 +18,8 @@ const mean = arr => arr.reduce((p, c) => p + c, 0) / arr.length;
 
 tables = $(".table-div");
 
+$("#extra-div").append("<div id='shift-diffs'></div>");
+
 
 function generateExtra() {
     $(".auto-generated").remove();
@@ -158,7 +160,7 @@ function generateExtra() {
 
             bank_holidays_worked = bank_holidays.length;
 
-            nights_worked= $(tr).find(".night-shift").length;
+            nights_worked = $(tr).find(".night-shift").length;
 
             shift_diff = JSON.parse(worker_td.attr("data-shift-diff"));
             summed_shift_diff = Object.values(shift_diff).reduce((a, b) => a + b);
@@ -166,6 +168,8 @@ function generateExtra() {
             total_summed_shift_diffs = total_summed_shift_diffs + summed_shift_diff;
             total_summed_shift_diffs_abs = total_summed_shift_diffs_abs + Math.abs(summed_shift_diff);
 
+
+
             worker_td.get(0).dataset.shiftDiffSummed = summed_shift_diff.toPrecision(2);
 
 
@@ -237,8 +241,10 @@ function generateExtra() {
         });
 
 
-        $(table).before($(`<span>Total summed shift diff: ${total_summed_shift_diffs}</span><br/>`));
+            $("#shift-diffs").before($(`<span>Total summed shift diff: ${total_summed_shift_diffs}</span><br/>`));
-        $(table).before($(`<span>Total summed shift diff: ${total_summed_shift_diffs_abs}</span><br/>`));
+            $("#shift-diffs").before($(`<span>Total summed shift diff (abs): ${total_summed_shift_diffs_abs}</span><br/>`));
+
+
         $(table).before(summary_button);
 
 
@@ -295,10 +301,8 @@ $(".table-div .worker-row .worker").each((n, tr) => {
     locum_shift_counts = jtr.data("locum-shift-counts")
     shift_targets = jtr.data("worker-targets")
 
-    console.log(shift_counts, locum_shift_counts, shift_targets)
 
     oshifts.forEach((s) => {
-        console.log(s, shift_counts, shift_targets, locum_shift_counts)
         if (s in shift_targets && shift_targets[s] > 0) {
             if (s in shift_counts) {
                 c = shift_counts[s];
@@ -309,7 +313,7 @@ $(".table-div .worker-row .worker").each((n, tr) => {
             if (s in locum_shift_counts) {
                 locum_shifts = `+${locum_shift_counts[s]}`
             }
-            row.append(`<td class="target-assigned ${s}" data-assigned=${c} data-target=${shift_targets[s].toFixed(2)} data-diff=${shift_targets[s].toFixed(2)-c}>${c}${locum_shifts} (${shift_targets[s].toFixed(2)})</td>`)
+            row.append(`<td class="target-assigned ${s}" data-assigned=${c} data-target=${shift_targets[s].toFixed(2)} data-diff=${shift_targets[s].toFixed(2) - c}>${c}${locum_shifts} (${shift_targets[s].toFixed(2)})</td>`)
 
         } else {
             row.append(`<td>-</td>`)
@@ -364,33 +368,454 @@ oshifts.forEach((shift) => {
 
 });
 
-oshifts.forEach((shift) => {
+let selectedShifts = new Set();
-    button = $(`<button data-shift='${shift}'>${shift}</button>`)
 
-    button.on("click", (evt) => {
+let shiftColors = {};
-        console.log(evt.target.dataset.shift)
-        $("#shift-timetable-div").empty();
-        table = $("<table>");
-        $("table#main-table th.date").each((n, th) => {
-            row = $("<tr>")
-            row.append(`<td>${th.dataset.date}</td>`);
-            $(`table#main-table td[data-date='${th.dataset.date}']`).filter(function() {
-                let shifts = $(this).attr("data-shift");
-                if (!shifts) return false;
-                return shifts.split(",").map(s => s.trim()).includes(evt.target.dataset.shift);
-            }).each((n, td) => {
-                row.append($(td).closest("tr").children("td:first").clone());
-            });
 
-            table.append(row)
+oshifts.forEach((shift, i) => {
-        });
+    selectedShifts.add(shift);
-        $("#shift-timetable-div").append(table);
+    // Assign a unique color for each shift using HSL
-
+    shiftColors[shift] = `hsl(${(i * 360 / oshifts.length)}, 70%, 80%)`;
-    })
-
-    $("#shift-timetable-options").append(button)
 });
 
+// Add a button to toggle grid view
+$("#shift-timetable-options").append(
+    `<button id="toggle-grid-view" style="margin-left:10px;">Toggle Linear/Grid View</button><br/>`
+);
+
+let gridViewEnabled = false;
+
+$("#toggle-grid-view").on("click", function () {
+    gridViewEnabled = !gridViewEnabled;
+    renderShiftTimetable();
+});
+
+
+const shiftColourControls = $(`
+    <details><summary>Shift Timetable Colours</summary>
+    <div id="shift-colour-controls" style="margin:16px 0;">
+        <b>Shift timetable colours:</b>
+        <span id="shift-colour-buttons"></span>
+    </div>
+    </details>
+`);
+$("#shift-timetable-options").append(shiftColourControls);
+
+// Add colour pickers for each shift
+oshifts.forEach(shift => {
+    // Colour picker for the shift
+    let colorInput = $(`<input type="color" class="shift-colour-picker" data-shift="${shift}" value="${rgb2hex(shiftColors[shift])}" style="margin-left:2px; vertical-align:middle;" title="Pick colour for ${shift}">`);
+    // Label for the shift
+    let label = $(`<label style="margin-right:10px;">${shift}</label>`);
+    // On change, update the shiftColors map, button color, and re-render the timetable
+    colorInput.on("input", function () {
+        shiftColors[shift] = $(this).val();
+        // Update the corresponding shift button color
+        $(`#shift-timetable-options button[data-shift='${shift}']`).css("background", selectedShifts.has(shift) ? shiftColors[shift] : "");
+        renderShiftTimetable();
+    });
+    $("#shift-colour-buttons").append(label.append(colorInput));
+});
+// Add the toggle button if not already present
+if ($("#show-all-shifts").length === 0) {
+    $("#shift-timetable-options").append(
+        `<button id="show-all-shifts" style="margin-left:10px;">Show All Shifts</button>`
+    );
+}
+
+let allShiftsShown = false; // Start with all shifts shown by default
+
+function updateShowAllShiftsButton() {
+    if (allShiftsShown) {
+        $("#show-all-shifts").addClass("active").text("Hide All Shifts");
+    } else {
+        $("#show-all-shifts").removeClass("active").text("Show All Shifts");
+    }
+}
+
+// Toggle logic
+$("#show-all-shifts").off("click").on("click", function () {
+    allShiftsShown = !allShiftsShown;
+    if (allShiftsShown) {
+        // Select all shifts
+        oshifts.forEach((shift) => {
+            selectedShifts.add(shift);
+            $(`#shift-timetable-options button[data-shift='${shift}']`)
+                .addClass("selected")
+                .css("background", shiftColors[shift]);
+        });
+    } else {
+        // Deselect all shifts
+        oshifts.forEach((shift) => {
+            selectedShifts.delete(shift);
+            $(`#shift-timetable-options button[data-shift='${shift}']`)
+                .removeClass("selected")
+                .css("background", "");
+        });
+    }
+    updateShowAllShiftsButton();
+    renderShiftTimetable();
+});
+
+// Ensure the button state is correct on page load and after timetable render
+updateShowAllShiftsButton();
+// Utility to convert hsl to hex for initial value
+function rgb2hex(rgb) {
+    // Accepts hsl() or rgb() or hex
+    if (rgb.startsWith("#")) return rgb;
+    if (rgb.startsWith("hsl")) {
+        // Convert hsl to rgb
+        let hsl = rgb.match(/hsl\(([\d.]+),\s*([\d.]+)%,\s*([\d.]+)%\)/);
+        if (!hsl) return "#ffe066";
+        let h = Number(hsl[1]), s = Number(hsl[2]) / 100, l = Number(hsl[3]) / 100;
+        let a = s * Math.min(l, 1 - l);
+        function f(n) {
+            let k = (n + h / 30) % 12;
+            let color = l - a * Math.max(Math.min(k - 3, 9 - k, 1), -1);
+            return Math.round(255 * color);
+        }
+        return "#" + [f(0), f(8), f(4)].map(x => x.toString(16).padStart(2, "0")).join("");
+    }
+    if (rgb.startsWith("rgb")) {
+        let rgbArr = rgb.match(/\d+/g).map(Number);
+        return "#" + rgbArr.map(x => x.toString(16).padStart(2, "0")).join("");
+    }
+    return "#ffe066";
+}
+
+// --- Add worker checkboxes to the shift timetable section ---
+function renderWorkerCheckboxes() {
+    // Get all unique worker names and IDs from the main table
+    let workerList = [];
+    $("table#main-table td.worker").each(function () {
+        let workerId = $(this).attr("data-worker-id");
+        let workerName = $(this).find("span.name").text().trim();
+        if (workerId && !workerList.some(w => w.id === workerId)) {
+            workerList.push({ id: workerId, name: workerName });
+        }
+    });
+
+    // Build checkboxes
+    let html = `<details id="shift-timetable-worker-filter" style="margin-bottom:8px;">
+        <summary><b>Show workers:</b></summary>`;
+    // Add a "Toggle All" button
+    html += `<button type="button" id="toggle-all-workers" style="margin-right:10px;">Toggle All</button>`;
+
+    workerList.forEach(w => {
+        html += `<label style="margin-right:8px;">
+            <input type="checkbox" name="worker-filter-checkbox" value="${w.id}" checked> ${w.name}
+        </label>`;
+    });
+
+    // Attach toggle logic after DOM insertion
+    setTimeout(() => {
+        $("#toggle-all-workers").off("click").on("click", function() {
+            const $checkboxes = $("input[name='worker-filter-checkbox']");
+            const allChecked = $checkboxes.length === $checkboxes.filter(":checked").length;
+            $checkboxes.prop("checked", !allChecked);
+            // Update the filter set and re-render once, instead of triggering "change" for each
+            workerFilterSet = new Set(
+                Array.from(document.querySelectorAll("input[name='worker-filter-checkbox']:checked"))
+                    .map(cb => cb.value)
+            );
+            // Call your function here after the map has finished
+            renderShiftTimetable();
+        });
+    }, 0);
+    html += `</details>`;
+
+    // Prepend to the shift timetable options
+    $("#shift-timetable-options").prepend(html);
+}
+
+// Store the current filter value (set of workerIds)
+let workerFilterSet = new Set();
+// Store the current setting
+let showShiftNames = false;
+
+// Call this after the DOM is ready and after timetable is rendered
+renderWorkerCheckboxes();
+
+// On first render, select all workers by default
+$("input[name='worker-filter-checkbox']").each(function() {
+    workerFilterSet.add($(this).val());
+});
+renderShiftTimetable();
+
+// Update the filter set on checkbox change
+$(document).on("change", "input[name='worker-filter-checkbox']", function() {
+    workerFilterSet = new Set(
+        $("input[name='worker-filter-checkbox']:checked").map(function() { return $(this).val(); }).get()
+    );
+    renderShiftTimetable();
+});
+
+let showWeekStart = false;
+// Add a checkbox to show/hide shift names in the shift timetable section
+$("#shift-timetable-options").prepend(`
+    <div id="shift-timetable-show-shift-names" style="margin-bottom:8px;">
+        <label>
+            <input type="checkbox" id="show-shift-names-checkbox">
+            Show shift names in timetable
+        </label>
+    </div>
+`).prepend(`
+    <div id="week-start-control" style="margin:8px 0;">
+        <label>
+            <input type="checkbox" id="show-week-start-checkbox">
+            Show week start dates in grid view
+        </label>
+    </div>
+`);
+
+// Add a button to remove all shift timetable colours (set to white)
+$("#shift-colour-controls").append(
+    `<button id="remove-shift-colours" style="margin-left:10px;">Remove Shift Colours</button>`
+);
+
+$("#remove-shift-colours").on("click", function () {
+    // Set all shift colours to white
+    oshifts.forEach(shift => {
+        shiftColors[shift] = "#ffffff";
+        // Update the colour picker if present
+        $(`.shift-colour-picker[data-shift='${shift}']`).val("#ffffff");
+        // Update the corresponding shift button color if selected
+        $(`#shift-timetable-options button[data-shift='${shift}']`).css("background", selectedShifts.has(shift) ? "#ffffff" : "");
+    });
+    renderShiftTimetable();
+});
+
+$("#show-shift-names-checkbox").on("change", function() {
+    showShiftNames = $(this).is(":checked");
+    renderShiftTimetable();
+});
+
+
+$("#show-week-start-checkbox").on("change", function () {
+    showWeekStart = $(this).is(":checked");
+    renderShiftTimetable();
+});
+
+// Patch renderShiftTimetable to show shift names if enabled
+function renderShiftTimetable() {
+    console.log("Render shift timetable");
+    console.log(workerFilterSet)
+    $("#shift-timetable-div").empty();
+    if (selectedShifts.size === 0) return;
+    if (workerFilterSet.size === 0) return;
+
+    // Helper: returns true if worker id matches filter set or set is empty (all)
+    function workerMatchesFilter(workerId) {
+        if (workerFilterSet.size === 0) return true;
+        return workerFilterSet.has(workerId);
+    }
+
+    if (!gridViewEnabled) {
+        // Linear view
+        let table = $("<table id='linear-shift-timetable'>");
+        let headerRow = $("<tr>");
+        headerRow.append("<th>Date</th><th>Week</th><th>Day</th><th>Worker(s)</th>");
+        table.append(headerRow);
+
+        $("table#main-table th.date").each((n, th) => {
+            let date = th.dataset.date;
+            let week = "";
+            let day = "";
+            let firstTd = $(`table#main-table td[data-date='${date}']`).first();
+            if (firstTd.length) {
+                week = firstTd.attr("data-week") || "";
+                day = firstTd.attr("data-day") || "";
+            }
+
+            let row = $("<tr>");
+            row.append(`<td>${date}</td>`);
+            row.append(`<td>${week}</td>`);
+            row.append(`<td>${day}</td>`);
+
+            let workerCells = [];
+            $(`table#main-table td[data-date='${date}']`).filter(function () {
+                let shifts = $(this).attr("data-shift");
+                if (!shifts) return false;
+                let cell_shifts = shifts.split(",").map(s => s.trim());
+                return Array.from(selectedShifts).some(s => cell_shifts.includes(s));
+            }).each((n, td) => {
+                let workerTd = $(td).closest("tr").children("td:first");
+                let workerName = workerTd.find("span.name").text().trim();
+                let workerId = workerTd.attr("data-worker-id");
+                if (!workerMatchesFilter(workerId)) return;
+                let cell_shifts = $(td).attr("data-shift").split(",").map(s => s.trim());
+                let matchedShifts = Array.from(selectedShifts).filter(s => cell_shifts.includes(s));
+                let color = "#fff";
+                let shiftLabel = "";
+                if (matchedShifts.length === 1) {
+                    color = shiftColors[matchedShifts[0]];
+                    if (showShiftNames) shiftLabel = `<span style="font-size:0.9em; color:#333;"> (${matchedShifts[0]})</span>`;
+                } else if (matchedShifts.length > 1) {
+                    color = `linear-gradient(90deg, ${shiftColors[matchedShifts[0]]} 0 50%, ${shiftColors[matchedShifts[1]]} 50% 100%)`;
+                    if (showShiftNames) shiftLabel = `<span style="font-size:0.9em; color:#333;"> (${matchedShifts.join(", ")})</span>`;
+                }
+                workerCells.push(`<td style="background:${color}">${workerName}${shiftLabel}</td>`);
+            });
+
+            if (workerCells.length === 0) {
+                row.append("<td></td>");
+            } else {
+                row.append(workerCells.join(""));
+            }
+
+            table.append(row);
+        });
+        $("#shift-timetable-div").append(table);
+        return;
+    }
+
+    // Grid view
+    let weeks = [];
+    let daysOfWeek = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"];
+    $("table#main-table td[data-week]").each(function () {
+        let week = $(this).attr("data-week");
+        if (week && !weeks.includes(week)) weeks.push(week);
+    });
+    weeks = weeks.sort((a, b) => parseInt(a) - parseInt(b));
+
+
+    // Build week -> start date map (earliest date in that week)
+    let weekStartDates = {};
+    $("table#main-table td[data-week][data-date]").each(function () {
+        let wk = $(this).attr("data-week");
+        let d = $(this).attr("data-date");
+        if (!wk || !d) return;
+        if (!weekStartDates[wk] || new Date(d) < new Date(weekStartDates[wk])) {
+            weekStartDates[wk] = d;
+        }
+    });
+
+    let workerIds = [];
+    let workerIdToName = {};
+    $("table#main-table td.worker").each(function () {
+        let workerId = $(this).attr("data-worker-id");
+        let workerName = $(this).find("span.name").text().trim();
+        if (workerId && !workerIds.includes(workerId)) {
+            if (workerMatchesFilter(workerId)) {
+                workerIds.push(workerId);
+                workerIdToName[workerId] = workerName;
+            }
+        }
+    });
+
+    let cellMap = {};
+    $("table#main-table td.worker").each(function () {
+        let workerId = $(this).attr("data-worker-id");
+        let workerName = $(this).find("span.name").text().trim();
+        if (!cellMap[workerId]) cellMap[workerId] = { name: workerName, cells: {} };
+    });
+    $("table#main-table td[data-week][data-day]").each(function () {
+        let td = $(this);
+        let tr = td.closest("tr");
+        let workerTd = tr.find("td.worker");
+        let workerId = workerTd.attr("data-worker-id");
+        let week = td.attr("data-week");
+        let day = td.attr("data-day");
+        let shifts = (td.attr("data-shift") || "").split(",").map(s => s.trim()).filter(Boolean);
+        if (workerId && week && day) {
+            if (!cellMap[workerId].cells[week]) cellMap[workerId].cells[week] = {};
+            cellMap[workerId].cells[week][day] = shifts;
+        }
+    });
+
+    let gridTable = $(`<table class="shift-grid-table" style="margin-bottom:20px; border-collapse:collapse; width:auto;"><caption>Shift Grid</caption></table>`);
+    let header = $("<tr><th style='border:1px solid #888; min-width:60px; width:80px;'>Week</th></tr>");
+    daysOfWeek.forEach(day => header.append(`<th style='border:1px solid #888;'>${day}</th>`));
+    gridTable.append(header);
+
+    weeks.forEach(week => {
+        let row = $(`<tr></tr>`);
+        // Show week number and (optionally) the week's start date
+        let weekLabel = [week];
+        if (showWeekStart && weekStartDates[week]) {
+            weekLabel += `<div style="font-size:0.85em; color:#444;">${weekStartDates[week]}</div>`;
+        }
+        row.append(`<td style="border:1px solid #888; min-width:60px; width:120px; font-weight:bold; vertical-align:middle;">${weekLabel}</td>`);
+        daysOfWeek.forEach(day => {
+            let cellContent = [];
+            for (const workerId of workerIds) {
+                let worker = cellMap[workerId];
+                let shifts = (worker.cells[week] && worker.cells[week][day]) || [];
+                let matchedShifts = Array.from(selectedShifts).filter(s => shifts.includes(s));
+                if (matchedShifts.length > 0) {
+                    let color = "#fff";
+                    let shiftLabel = "";
+                    if (matchedShifts.length === 1) {
+                        color = shiftColors[matchedShifts[0]];
+                        if (showShiftNames) shiftLabel = `<span style="font-size:0.9em; color:#333;"> (${matchedShifts[0]})</span>`;
+                    } else if (matchedShifts.length > 1) {
+                        color = `linear-gradient(90deg, ${shiftColors[matchedShifts[0]]} 0 50%, ${shiftColors[matchedShifts[1]]} 50% 100%)`;
+                        if (showShiftNames) shiftLabel = `<span style="font-size:0.9em; color:#333;"> (${matchedShifts.join(", ")})</span>`;
+                    }
+                    cellContent.push(`<div style="background:${color};margin:1px 0;padding:0 2px;border-radius:3px;display:inline-block;">${worker.name}${shiftLabel}</div>`);
+                }
+            }
+            row.append(`<td style="border:1px solid #888; min-width:40px; text-align:center;">${cellContent.join(" / ")}</td>`);
+        });
+        gridTable.append(row);
+    });
+    $("#shift-timetable-div").append(gridTable);
+}
+
+// Re-render worker checkboxes after timetable is updated (in case workers change)
+function rerenderWorkerCheckboxesIfNeeded() {
+    $("#shift-timetable-worker-filter").remove();
+    renderWorkerCheckboxes();
+    // Restore checked state
+    $("input[name='worker-filter-checkbox']").each(function() {
+        if (workerFilterSet.has($(this).val())) {
+            $(this).prop("checked", true);
+        }
+    });
+}
+
+// Call this after timetable is (re)rendered if needed
+rerenderWorkerCheckboxesIfNeeded();
+
+// Patch the timetable button logic to use renderShiftTimetable
+oshifts.forEach((shift) => {
+    let button = $(`<button data-shift='${shift}'>${shift}</button>`);
+    button.css({
+        background: "",
+        border: "1px solid #888",
+        margin: "2px",
+        color: "#222"
+    });
+
+    button.on("click", function (evt) {
+        let s = evt.target.dataset.shift;
+        if (selectedShifts.has(s)) {
+            selectedShifts.delete(s);
+            $(this).removeClass("selected");
+            $(this).css("background", "");
+        } else {
+            selectedShifts.add(s);
+            $(this).addClass("selected");
+            $(this).css("background", shiftColors[shift]);
+        }
+        renderShiftTimetable();
+    });
+
+    $("#shift-timetable-options").append(button);
+});
+
+
+// Add help text to the shift timetable section
+$("#shift-timetable-options").prepend(`
+    <details id="shift-timetable-help" style="margin-bottom:8px; font-size: 0.95em; color: #444;">
+        <summary><b>Shift Timetables Help:</b></summary>
+        <ul style="margin: 4px 0 0 18px; padding: 0;">
+            <li>Select one or more shifts below to view which workers are assigned on each day.</li>
+            <li>Use the colour pickers to customise shift colours in the timetable and main table.</li>
+            <li>Toggle between linear and grid view using the "Toggle Linear/Grid View" button.</li>
+        </ul>
+    </details>
+`);
 //$("table.summary")
 
 // $("body").prepend("<div id='global-settings'>Settings:<br />Customise training days lost per shift (you need to press recalculate for changes to take effect)<form></form></div>")
@@ -448,13 +873,13 @@ $("td").hover((e) => {
     if (targetSpan) {
         // Only highlight spans with the same text
         let shiftDisplay = $(targetSpan).text().trim();
-        $("span.multi-shift-shift").filter(function() {
+        $("span.multi-shift-shift").filter(function () {
             return $(this).text().trim() === shiftDisplay;
         }).addClass("shift-highlight");
     } else if (e.target.dataset.shift != undefined && e.target.dataset.shift.length > 0) {
         // Fallback: highlight all cells containing ANY of the hovered shifts (legacy)
         let hovered_shifts = e.target.dataset.shift.split(",").map(s => s.trim());
-        $("td").filter(function() {
+        $("td").filter(function () {
             let shifts = $(this).attr("data-shift");
             if (!shifts) return false;
             let cell_shifts = shifts.split(",").map(s => s.trim());
@@ -605,3 +1030,182 @@ function syntaxHighlight(json) {
         return '<span class="' + cls + '">' + match + '</span>';
     });
 }
+
+(function() {
+    const mainTable = document.getElementById("main-table");
+    if (!mainTable) return;
+
+    // Helper to remove highlight from all headers
+    function clearColHover() {
+        const ths = mainTable.querySelectorAll("th.col-hover");
+        ths.forEach(th => th.classList.remove("col-hover"));
+    }
+
+    // Mouseover handler for all cells
+    mainTable.addEventListener("mouseover", function(e) {
+        let cell = e.target.closest("td,th");
+        if (!cell) return;
+        // Find the column index
+        let colIdx = cell.cellIndex;
+        if (colIdx === undefined) return;
+        clearColHover();
+        // Highlight the header in this column (first row)
+        let headerRow = mainTable.tHead ? mainTable.tHead.rows[0] : mainTable.rows[0];
+        if (headerRow && headerRow.cells[colIdx]) {
+            headerRow.cells[colIdx].classList.add("col-hover");
+        }
+    });
+
+    // Mouseout handler to clear highlight
+    mainTable.addEventListener("mouseout", function(e) {
+        clearColHover();
+    });
+
+
+})();
+
+// --- Add a button to show the Display Settings modal ---
+$("#main-table").before(
+    `<button id="show-display-settings" style="margin:8px 0 8px 0; float:right;">Display Settings</button>`
+);
+
+// --- Modal show/hide logic ---
+$("#show-display-settings").on("click", function() {
+    $("#display-settings-modal").show();
+});
+// Attach the event handler after the modal is added to the DOM
+$(document).on("click", "#close-display-settings", function() {
+    $("#display-settings-modal").hide();
+});
+
+// --- Modal HTML (replace the highlight-days-controls div with this) ---
+const displaySettingsModal = $(`
+<div id="display-settings-modal" style="display:none; position:fixed; top:50%; left:50%; transform:translate(-50%,-50%); z-index:10000; background:#fff; border:2px solid #888; border-radius:8px; box-shadow:0 4px 24px #0002; padding:24px; min-width:320px;">
+    <h3 style="margin-top:0;">Display Settings</h3>
+    <div id="highlight-days-controls" style="margin:8px 0;">
+        <b>Highlight days of week:</b>
+        <span id="highlight-days-buttons"></span>
+    </div>
+    <div style="text-align:right; margin-top:16px;">
+        <button id="close-display-settings">Close</button>
+    </div>
+</div>
+`);
+$("body").append(displaySettingsModal);
+
+// --- Highlight day buttons in modal with color pickers ---
+const daysOfWeek = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"];
+let highlightedDays = new Set();
+let highlightDayColors = {}; // day -> color
+
+// Default highlight color
+const defaultHighlightColor = "#ffe066";
+
+// --- Add force assigned highlight controls to the Display Settings modal ---
+const forceAssignedControls = $(`
+    <div id="force-assigned-controls" style="margin:16px 0;">
+        <label>
+            <input type="checkbox" id="highlight-force-assigned" checked>
+            Highlight force assigned shifts
+        </label>
+    </div>
+`);
+$("#display-settings-modal #highlight-days-controls").after(forceAssignedControls);
+
+// --- Highlighting logic ---
+function updateDayHighlights() {
+    // Remove previous highlights and reset borders/backgrounds
+    $("#main-table td, #main-table th").removeClass("day-highlighted")
+        .css("border-left", "").css("border-right", "").css("background", "");
+
+    // For each selected day, highlight all cells in that column
+    highlightedDays.forEach(day => {
+        let color = highlightDayColors[day] || defaultHighlightColor;
+        // Highlight th: only change border color, not background
+        $(`#main-table th[data-day='${day}']`)
+            .addClass("day-highlighted")
+            .css("border-left", `3px solid ${color}`)
+            .css("border-right", `3px solid ${color}`)
+            .css("background", "");
+        // Highlight td: only change border color, not background
+        $(`#main-table td[data-day='${day}']`)
+            .addClass("day-highlighted")
+            .css("border-left", `3px solid ${color}`)
+            .css("border-right", `3px solid ${color}`)
+            .css("background", "");
+    });
+    // Update modal button states
+    $("#highlight-days-buttons .highlight-day-btn").each(function() {
+        let day = $(this).data("day");
+        if (highlightedDays.has(day)) {
+            $(this).addClass("highlighted");
+        } else {
+            $(this).removeClass("highlighted");
+        }
+    });
+}
+
+// --- Highlighting logic for force assigned shifts ---
+function updateForceAssignedHighlights() {
+    if ($("#highlight-force-assigned").prop("checked")) {
+        // Add highlight to all force assigned cells
+        $("#main-table td.force-assigned").addClass("force-assigned-highlight");
+    } else {
+        // Remove highlight
+        $("#main-table td.force-assigned").removeClass("force-assigned-highlight");
+    }
+}
+
+// --- Update both day and force assigned highlights together ---
+function updateAllHighlights() {
+    updateDayHighlights();
+    updateForceAssignedHighlights();
+}
+
+// --- Update force assigned highlights when checkbox changes ---
+$("#display-settings-modal").on("change", "#highlight-force-assigned", function() {
+    updateForceAssignedHighlights();
+});
+
+
+
+
+// --- Also call updateAllHighlights on modal open ---
+$("#show-display-settings").on("click", function() {
+    $("#display-settings-modal").show();
+    updateAllHighlights();
+});
+
+// --- Highlight day buttons in modal with color pickers ---
+daysOfWeek.forEach(day => {
+    // Button to toggle highlight
+    let btn = $(`<button type="button" class="highlight-day-btn" data-day="${day}" style="margin:0 2px;">${day}</button>`);
+    // Color picker
+    let colorInput = $(`<input type="color" class="highlight-day-color" data-day="${day}" value="${defaultHighlightColor}" style="margin-left:2px; vertical-align:middle;" title="Pick highlight color for ${day}">`);
+    highlightDayColors[day] = defaultHighlightColor;
+
+    btn.on("click", function () {
+        if (highlightedDays.has(day)) {
+            highlightedDays.delete(day);
+            $(this).removeClass("highlighted");
+        } else {
+            highlightedDays.add(day);
+            $(this).addClass("highlighted");
+        }
+        updateAllHighlights();
+    });
+
+    colorInput.on("input", function () {
+        highlightDayColors[day] = $(this).val();
+        updateAllHighlights();
+    });
+
+    $("#highlight-days-buttons").append(btn).append(colorInput);
+});
+
+
+// --- Initial highlight on page load ---
+$(function() {
+    updateAllHighlights();
+
+});

rota/workers.py

@@ -1,7 +1,7 @@
 import datetime
 from collections import defaultdict
-from typing import Iterable, List, Literal, Optional
+from typing import Iterable, List, Literal, Optional, Set
-from pydantic import BaseModel, ConfigDict, field_validator
+from pydantic import BaseModel, ConfigDict, field_validator, Field
 
 from rich.pretty import pprint
 from rota.console import console
@@ -31,6 +31,31 @@ class NotAvailableToWork(BaseModel):
     date: datetime.date
     reason: str = "unknown"
 
+class HardDayDependency(BaseModel):
+    days: Set[str] = Field(..., description="List of days for the dependency group")
+    weeks: Optional[List[int]] = None  # If None, applies to all weeks
+    start_date: Optional[datetime.date] = None  # Optional start date for the dependency
+    end_date: Optional[datetime.date] = None  # Optional end date for the dependency
+
+    @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 HardDayExclusion(BaseModel):
+    days: set[str] = Field(..., description="List of days to exclude together")
+    weeks: Optional[List[int]] = None  # If None, applies to all weeks
+    start_date: Optional[datetime.date] = None  # Optional start date for the dependency
+    end_date: Optional[datetime.date] = None  # Optional end date for the dependency
+
+    @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
 
 def generate_not_available_to_works(
     start_date: datetime.date,
@@ -114,10 +139,13 @@ class Worker(BaseModel):
     allowed_multi_shift_sets: list[set] = []  # or list/set of frozensets
     prefer_multi_shift_together: int = 0  # Default: no preference
     # Set to a positive integer to prefer, negative to discourage, 0 for neutral
-    hard_day_dependencies: set[frozenset[str]] = set()  # e.g. {frozenset({"Mon", "Tue"}), frozenset({"Fri", "Sun"})}
+    hard_day_dependencies: list[HardDayDependency] = []  
-    hard_day_exclusions: list[tuple[str, str]] = []  # e.g. [("Fri", "Sun")]
+    hard_day_exclusions: list[HardDayExclusion] = []  
     force_assign_with: dict[str, list[str]] = {}  # e.g. {"oncall": ["weekend"]}
     max_shifts_per_week_by_shift_name: dict[str, int] = {}  # e.g. {"night": 2, "a": 3}
+    forced_assignments: List[tuple[int, str, str]] = []  # (week, day, shift_name)
+    forced_assignments_by_date: List[tuple[datetime.date, str]] = []  # (date, shift_name)
+    max_days_worked_per_week: int | None= None # Default: no limit, can be set to a specific number
 
 
     assign_as_block_preferences: dict[str, float] = {}  # shift_name -> weight (positive = prefer block)
@@ -409,15 +437,21 @@ class Worker(BaseModel):
             self.allowed_multi_shift_sets = []
         self.allowed_multi_shift_sets.append(frozenset(shifts))
 
-    def add_hard_day_dependency(self, *days: str):
+    def add_hard_day_dependency(self, days: set[str], weeks: Optional[List[int]] = None, start_date: Optional[datetime.date] = None, end_date: Optional[datetime.date] = None):
-        if not hasattr(self, "hard_day_dependencies"):
+        """
-            self.hard_day_dependencies = set()
+        Add a hard day dependency for this worker.
-        self.hard_day_dependencies.add(frozenset(days))
+        days: List of days that must be grouped together.
+        weeks: Optional list of weeks this dependency applies to. If None, applies to all weeks.
+        """
+        self.hard_day_dependencies.append(HardDayDependency(days=days, weeks=weeks, start_date=start_date, end_date=end_date))
 
-    def add_hard_day_exclusion(self, day1: str, day2: str):
+    def add_hard_day_exclusion(self, days: set[str], weeks: Optional[List[int]] = None, start_date: Optional[datetime.date] = None, end_date: Optional[datetime.date] = None):
-        if not hasattr(self, "hard_day_exclusions"):
+        """
-            self.hard_day_exclusions = []
+        Add a hard day exclusion for this worker.
-        self.hard_day_exclusions.append((day1, day2))
+        days: List of days that must be excluded together.
+        weeks: Optional list of weeks this exclusion applies to. If None, applies to all weeks.
+        """
+        self.hard_day_exclusions.append(HardDayExclusion(days=days, weeks=weeks, start_date=start_date, end_date=end_date))
 
     def add_force_assign_with(self, shift: str, with_shifts: list[str] | str):
         if not hasattr(self, "force_assign_with"):
@@ -431,3 +465,11 @@ class Worker(BaseModel):
         if not hasattr(self, "max_shifts_per_week_by_shift_name"):
             self.max_shifts_per_week_by_shift_name = {}
         self.max_shifts_per_week_by_shift_name[shift_name] = max_per_week
+
+    def force_assign_shift(self, week: int, day: str, shift_name: str):
+        """Force this worker to be assigned to a shift on a specific week/day."""
+        self.forced_assignments.append((week, day, shift_name))
+
+    def force_assign_shift_by_date(self, date: datetime.date, shift_name: str):
+        """Force this worker to be assigned to a shift on a specific date."""
+        self.forced_assignments_by_date.append((date, shift_name))

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_shifts.py

@@ -1,717 +1,444 @@
 import datetime
 import pytest
-from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, WarningTermination, days, WorkerRequirement
+from rota.shifts import (
-
+    InvalidShift, MaxShiftsPerWeekBlockConstraint, MaxShiftsPerWeekConstraint,
+    NoWorkers, PostShiftConstraint, PreShiftConstraint, RotaBuilder, SingleShift,
+    WarningTermination, days, WorkerRequirement
+)
 from rota.workers import NotAvailableToWork, WorkRequests, Worker, generate_not_available_to_works
 
-import itertools
-
-
 def generate_basic_rota(weeks_to_rota=10, workers=2, start_date=datetime.date(2022, 3, 7)):
-
     Rota = RotaBuilder(
         start_date,
         weeks_to_rota=weeks_to_rota,
     )
-
-    # Add a few workers
     Rota.add_workers(
-        [
+        [Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)]
-            Worker(name=f"worker{i}", site="group1", grade=1) for i in range(1, workers+1)
-        ]
     )
-
     return Rota
 
-
 def date_generator(from_date, days):
     n = 0
     while True:
         yield from_date
-
+        n += 1
-        n = n + 1
         if n >= days:
             break
         from_date = from_date + datetime.timedelta(days=1)
 
-
+def test_duplicate_shifts():
-class TestShifts:
-    def test_duplicate_shifts(self):
     Rota = generate_basic_rota()
-
+    Rota.add_worker(Worker(name="worker3", site="group1", grade=1))
-        Rota.add_worker(
+    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days))
-            Worker(
-                name="worker3", site="group1", grade=1,
-            ),
-        )
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
-
+    Rota.add_shifts(SingleShift(sites=("group1",), name="a", length=12.5, days=days))
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1",), name="a", length= 12.5, days=days,
-            ),
-        )
-
     with pytest.raises(InvalidShift):
         Rota.build_and_solve(options={"ratio": 0.000})
 
-
+def test_max_shifts():
-class TestShiftConstraints:
-
-    def test_max_shifts(self):
     Rota = generate_basic_rota()
-
+    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-        Rota.add_shifts(
+                                constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]))
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
 
-    def test_max_shifts_fail(self):
+def test_max_shifts_fail():
     Rota = generate_basic_rota()
-
+    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-        Rota.add_shifts(
+                                constraints=[MaxShiftsPerWeekConstraint(max_shifts=3)]))
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
-                constraint=[{"name": "max_shifts_per_week", "options": 3}],
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-
     assert Rota.results.solver.termination_condition == "infeasible"
-
+    Rota.add_worker(Worker(name="worker3", site="group1", grade=1))
-        Rota.add_worker(
-            Worker(
-                name="worker3", site="group1", grade=1,
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
 
-
+def test_max_shifts_per_week_block():
-    def test_max_shifts_per_week_block(self):
     Rota = generate_basic_rota()
-
+    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-        Rota.add_shifts(
+                                constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2)]))
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
-                constraint=[{"name": "max_shifts_per_week_block", "options": {"weeks": 2}}],
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("max_shifts_per_week_block")
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
 
-
+def test_max_shifts_per_week_block_fail():
-
-    def test_max_shifts_per_week_block_fail(self):
     Rota = generate_basic_rota()
-
+    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-        Rota.add_shifts(
+                                constraints=[MaxShiftsPerWeekBlockConstraint(week_block=3)]))
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
-                constraint=[{"name": "max_shifts_per_week_block", "options": {"weeks": 3}}],
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("max_shifts_per_week_block")
-
     assert Rota.results.solver.status in ("warning", "error")
 
-    def test_max_shifts_per_week_block_shift_number(self):
+def test_max_shifts_per_week_block_shift_number():
     Rota = generate_basic_rota()
-
+    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-        Rota.add_shifts(
+                                constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2, max_shifts=7)]))
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
-                constraint=[{"name": "max_shifts_per_week_block", "options": {"weeks": 2, "shift_number": 7}}],
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("max_shifts_per_week_block")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_max_shifts_per_week_block_shift_number_fail(self):
+def test_max_shifts_per_week_block_shift_number_fail():
     Rota = generate_basic_rota()
-
+    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-        Rota.add_shifts(
+                                constraints=[MaxShiftsPerWeekBlockConstraint(week_block=2, max_shifts=6)]))
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
-                constraint=[{"name": "max_shifts_per_week_block", "options": {"weeks": 2, "shift_number": 6}}],
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("max_shifts_per_week_block")
-
     assert Rota.results.solver.status in ("warning", "error")
 
-    def test_pre_shift_constraint(self):
+def test_pre_shift_constraint():
     Rota = generate_basic_rota()
-
+    for i, d in enumerate(days[:6]):
-        Rota.add_shifts(
+        Rota.add_shifts(SingleShift(sites=("group1", "group2"), name=chr(97+i), length=12.5, days=d,
-            SingleShift(
+                                    constraints=[PreShiftConstraint(days=1)]))
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
-                constraint=[{"name": "pre", "options": 1}],
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="b", length= 12.5, days=days[1],
-                constraint=[{"name": "pre", "options": 1}],
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="c", length= 12.5, days=days[2],
-                constraint=[{"name": "pre", "options": 1}],
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="d", length= 12.5, days=days[3],
-                constraint=[{"name": "pre", "options": 1}],
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="e", length= 12.5, days=days[4],
-                constraint=[{"name": "pre", "options": 1}],
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="f", length= 12.5, days=days[5],
-                constraint=[{"name": "pre", "options": 1}],
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("pre")
-
     assert Rota.results.solver.status == "ok"
-
     for worker in Rota.get_workers():
         shift_string = Rota.get_worker_shift_list_string(worker)
-
-            print(shift_string)
         for s in ["ab", "bc", "cd", "de", "ef", "fg"]:
             assert s not in shift_string
-
+    Rota.add_shifts(SingleShift(sites=("group1", "group2"), name="g", length=12.5, days=days[6],
-            
+                                constraints=[PreShiftConstraint(days=2)]))
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"), name="g", length= 12.5, days=days[6],
-                constraint=[{"name": "pre", "options": 2}],
-            ),
-        )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-
     assert Rota.results.solver.status in ("warning", "error")
 
-    def test_pre_shift_constraint2(self):
+def test_pre_shift_constraint2():
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                    constraints=[PreShiftConstraint(days=1)], workers_required=2)
-                constraint=[{"name": "pre", "options": 2}],
-                workers_required=2
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
-                constraint=[{"name": "pre", "options": 3}],
-                workers_required=2
-            ),
-
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("pre")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_pre_shift_constraint3(self):
+def test_pre_shift_constraint3():
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                    constraints=[PreShiftConstraint(days=3)], workers_required=2),
-                constraint=[{"name": "pre", "options": 3}],
+        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
-                workers_required=2
+                    constraints=[PreShiftConstraint(days=3)], workers_required=2),
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
-                constraint=[{"name": "pre", "options": 3}],
-                workers_required=2
-            ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("pre")
-
     assert Rota.results.solver.status in ("warning", "error")
 
-    def test_post_shift_constraint(self):
+def test_post_shift_constraint():
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                    constraints=[PostShiftConstraint(days=3)], workers_required=2),
-                constraint=[{"name": "post", "options": 3}],
+        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
-                workers_required=2
+                    constraints=[PostShiftConstraint(days=2)], workers_required=2),
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
-                constraint=[{"name": "post", "options": 2}],
-                workers_required=2
-            ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("post")
-
     assert Rota.results.solver.status == "ok"
 
-    def test_post_shift_constraint2(self):
+def test_post_shift_constraint2():
     Rota = generate_basic_rota()
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[0],
+                    constraints=[PostShiftConstraint(days=3)], workers_required=2),
-                constraint=[{"name": "post", "options": 3}],
+        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4],
-                workers_required=2
+                    constraints=[PostShiftConstraint(days=3)], workers_required=2),
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="b", length= 12.5, days=days[4],
-                constraint=[{"name": "post", "options": 3}],
-                workers_required=2
-            ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("post")
-
     assert Rota.results.solver.status in ("warning", "error")
 
+def test_shift_constraint_week_limit():
+    Rota = generate_basic_rota()
+    Rota.add_shifts(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
+                    constraints=[PostShiftConstraint(days=1, weeks=[1,2,3])], workers_required=2),
+        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[1], workers_required=2),
+    )
+    Rota.build_and_solve(options={"ratio": 0.001})
+    Rota.export_rota_to_html("post")
+    assert Rota.results.solver.status in ("warning", "error")
+    Rota.add_workers([
+        Worker(name="worker3", end_date=Rota.start_date + datetime.timedelta(weeks=3), site="group1", grade=1),
+        Worker(name="worker4", end_date=Rota.start_date + datetime.timedelta(weeks=3), site="group1", grade=1)
+    ])
+    Rota.build_and_solve(options={"ratio": 0.001})
+    assert Rota.results.solver.status == "ok"
 
-class TestShiftDates:
+def test_shift_constraint_week_limit2():
+    Rota = generate_basic_rota()
+    Rota.add_shifts(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[0],
+                    constraints=[PostShiftConstraint(days=1, start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4))], workers_required=2),
+        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[1], workers_required=2),
+    )
+    Rota.build_and_solve(options={"ratio": 0.001})
+    assert Rota.results.solver.status in ("warning", "error")
+    Rota.add_workers([
+        Worker(name="worker3", start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4), site="group1", grade=1),
+        Worker(name="worker4", start_date=Rota.start_date + datetime.timedelta(weeks=2), end_date=Rota.start_date + datetime.timedelta(weeks=4), site="group1", grade=1)
+    ])
+    Rota.build_and_solve(options={"ratio": 0.001})
+    Rota.export_rota_to_html("post")
+    assert Rota.results.solver.status == "ok"
 
-    def test_shift_start_date(self):
+def test_shift_start_date():
     Rota = generate_basic_rota(workers=2)
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                    start_date=(Rota.start_date + datetime.timedelta(days=7))),
-                start_date=(Rota.start_date + datetime.timedelta(days=7)),
-            ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-        #Rota.export_rota_to_html("test_shift_start_date")
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
-
     assert Rota.get_workers_total_shifts()["worker1"] in (31, 32)
     assert Rota.get_workers_total_shifts()["worker2"] in (31, 32)
 
-    def test_shift_start_date_end_date(self):
+def test_shift_start_date_end_date():
     Rota = generate_basic_rota(workers=2)
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
                     start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
+                    end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
-            ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-        #Rota.export_rota_to_html("test_shift_start_date")
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
-
     assert Rota.get_workers_total_shifts()["worker1"] in (17, 18)
     assert Rota.get_workers_total_shifts()["worker2"] in (17, 18)
-
     for worker in Rota.get_workers():
         assert Rota.get_worker_shift_list_string(worker).startswith("-"*7)
         assert Rota.get_worker_shift_list_string(worker).endswith("-"*28)
 
-    def test_shift_invalid_start_date(self):
+def test_shift_invalid_start_date():
     Rota = generate_basic_rota(workers=2)
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                    start_date=(Rota.start_date - datetime.timedelta(days=6))),
-                start_date=(Rota.start_date - datetime.timedelta(days=6)),
-            ),
     )
-
     with pytest.raises(WarningTermination):
         Rota.build_and_solve(options={"ratio": 0.000})
 
-    def test_shift_invalid_start_date2(self):
+def test_shift_invalid_start_date2():
     Rota = generate_basic_rota(workers=2)
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                    start_date=(Rota.rota_end_date + datetime.timedelta(days=1))),
-                start_date=(Rota.rota_end_date + datetime.timedelta(days=1)),
-            ),
     )
-
     with pytest.raises(WarningTermination):
         Rota.build_and_solve(options={"ratio": 0.000})
 
-    def test_shift_invalid_end_date(self):
+def test_shift_invalid_end_date():
     Rota = generate_basic_rota(workers=2)
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                    end_date=(Rota.rota_end_date + datetime.timedelta(days=1))),
-                end_date=(Rota.rota_end_date + datetime.timedelta(days=1)),
-            ),
     )
-
     with pytest.raises(WarningTermination):
         Rota.build_and_solve(options={"ratio": 0.000})
 
-    def test_shift_invalid_end_date2(self):
+def test_shift_invalid_end_date2():
     Rota = generate_basic_rota(workers=2)
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                    end_date=(Rota.start_date - datetime.timedelta(days=1))),
-                end_date=(Rota.start_date - datetime.timedelta(days=1)),
-            ),
     )
-
     with pytest.raises(WarningTermination):
         Rota.build_and_solve(options={"ratio": 0.000})
 
-
+def test_shift_start_date_end_date_block():
-    def test_shift_start_date_end_date_block(self):
     Rota = generate_basic_rota(workers=2)
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days,
-                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                     assign_as_block=True,
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)],
                     start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
+                    end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
-            ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
-        #Rota.export_rota_to_html("test_shift_start_date")
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
-
     assert Rota.get_workers_total_shifts()["worker1"] in (17, 18)
     assert Rota.get_workers_total_shifts()["worker2"] in (17, 18)
-
     for worker in Rota.get_workers():
         assert Rota.get_worker_shift_list_string(worker).startswith("-"*7)
         assert Rota.get_worker_shift_list_string(worker).endswith("-"*28)
 
-    def test_shift_start_date_end_date_block_more_workers(self):
+def test_shift_start_date_end_date_block_more_workers():
     Rota = generate_basic_rota(workers=10)
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                    assign_as_block=True, force_as_block=True, workers_required=4,
-                assign_as_block=True,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]),
-                force_as_block=True,
+        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
-                workers_required=4,
+                    assign_as_block=True, force_as_block=True, workers_required=4,
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
-                #start_date=(Rota.start_date + datetime.timedelta(days=7)),
+                                 PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
-                #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="b", length= 12.5, days=days[4:],
-                assign_as_block=True,
-                force_as_block=True,
-                workers_required=4,
-                constraint=[{"name": "max_shifts_per_week", "options": 4},
-                             {"name": "pre", "options": 3}, {"name": "post", "options": 3}],
                     start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
+                    end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
-            ),
     )
-
     Rota.build_and_solve(options={"ratio": 0.000})
     Rota.export_rota_to_html("test_shift_start_date")
-
     assert Rota.results.solver.status == "ok"
     assert Rota.results.solver.termination_condition == "optimal"
-
     total_shifts = Rota.get_workers_total_shifts()
     for worker in Rota.get_workers():
-            print(Rota.get_worker_shift_list_string(worker))
         assert total_shifts[worker.name] == 22
-            assert "ab" not in Rota.get_worker_shift_list_string(worker)
-            assert "ba" not in Rota.get_worker_shift_list_string(worker)
-
-
-    def test_shift_start_date_end_date_block_with_requests(self):
-        Rota = generate_basic_rota(workers=7)
-
-
-        request_date = Rota.start_date + datetime.timedelta(days=(7*6)+4)
-        request = [WorkRequests(date=request_date, shift="c")]
-
-        Rota.add_workers(
-            [
-                Worker(name=f"worker-wr-{i}", site="group2", grade=1, work_requests=request) for i in range(1, 4)
-                #Worker(name="worker2", site="group1", grade=1),
-            ]
-        )
-
-        Rota.add_shifts(
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
-                assign_as_block=True,
-                force_as_block=True,
-                workers_required=4,
-                constraint=[{"name": "max_shifts_per_week", "options": 4}],
-                #start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="b", length= 12.5, days=days[4:],
-                assign_as_block=True,
-                force_as_block=True,
-                workers_required=4,
-                constraint=[{"name": "max_shifts_per_week", "options": 4},
-                             {"name": "pre", "options": 3}, {"name": "post", "options": 3}],
-                start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="c", length= 12.5, days=days[4:],
-                assign_as_block=True,
-                force_as_block=True,
-                workers_required=4,
-                constraint=[{"name": "max_shifts_per_week", "options": 4},
-                             {"name": "pre", "options": 3}, {"name": "post", "options": 3}],
-                start_date=(Rota.start_date + datetime.timedelta(days=(7*6))),
-            ),
-        )
-
-        Rota.build_and_solve(options={"ratio": 0.1})
-        Rota.export_rota_to_html("test_shift_start_date")
-
-        assert Rota.results.solver.status == "ok"
-        assert Rota.results.solver.termination_condition == "optimal"
-
-        total_shifts = Rota.get_workers_total_shifts()
-        for worker in Rota.get_workers():
-            print(Rota.get_worker_shift_list_string(worker))
-            assert total_shifts[worker.name] in range(24,30)
         shift_string = Rota.get_worker_shift_list_string(worker)
         assert "ab" not in shift_string
         assert "ba" not in shift_string
 
-            # Check all has been assigned as block
+def test_shift_start_date_end_date_block_with_requests():
+    Rota = generate_basic_rota(workers=7)
+    request_date = Rota.start_date + datetime.timedelta(days=(7*6)+4)
+    request = [WorkRequests(date=request_date, shift="c")]
+    Rota.add_workers([
+        Worker(name=f"worker-wr-{i}", site="group2", grade=1, work_requests=request) for i in range(1, 4)
+    ])
+    Rota.add_shifts(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
+                    assign_as_block=True, force_as_block=True, workers_required=4,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4)]),
+        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
+                    assign_as_block=True, force_as_block=True, workers_required=4,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
+                                 PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
+                    start_date=(Rota.start_date + datetime.timedelta(days=7)),
+                    end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1)))),
+        SingleShift(sites=("group1", "group2"), name="c", length=12.5, days=days[4:],
+                    assign_as_block=True, force_as_block=True, workers_required=4,
+                    constraints=[MaxShiftsPerWeekConstraint(max_shifts=4),
+                                 PreShiftConstraint(days=3), PostShiftConstraint(days=3)],
+                    start_date=(Rota.start_date + datetime.timedelta(days=(7*6)))),
+    )
+    Rota.build_and_solve(options={"ratio": 0.1})
+    Rota.export_rota_to_html("test_shift_start_date")
+    assert Rota.results.solver.status == "ok"
+    assert Rota.results.solver.termination_condition == "optimal"
+    total_shifts = Rota.get_workers_total_shifts()
+    for worker in Rota.get_workers():
+        assert total_shifts[worker.name] in range(24, 30)
+        shift_string = Rota.get_worker_shift_list_string(worker)
+        assert "ab" not in shift_string
+        assert "ba" not in shift_string
         assert shift_string.count("aaaa") == shift_string.count("a") / 4
         assert shift_string.count("bbb") == shift_string.count("b") / 3
         assert shift_string.count("ccc") == shift_string.count("c") / 3
-
-        # TODO test the request assignment
     for worker in Rota.get_workers_by_group()["group2"]:
         assert Rota.get_worker_shifts_by_date(worker)[request_date] == "c"
 
-class TestShiftWorkerRequirements:
+def test_worker_requirement_start_and_end_date():
-    def test_start_and_end_date(self):
     WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14))
 
-    def test_end_date_before_start_date(self):
+def test_worker_requirement_end_date_before_start_date():
     with pytest.raises(ValueError):
         WorkerRequirement(start_date=datetime.date(2022, 3, 14), end_date=datetime.date(2022, 3, 7))
 
-    def test_worker_requirement(self):
+def test_worker_requirement():
     Rota = generate_basic_rota(workers=2)
-
     wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21))]
-        print(wr)
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                    assign_as_block=True, force_as_block=True, workers_required=wr),
-                assign_as_block=True,
-                force_as_block=True,
-                workers_required=wr,
-                #start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
     )
     Rota.build_and_solve(options={"ratio": 0.1})
     Rota.export_rota_to_html("test_worker_requirement")
-
     assert Rota.results.solver.status == "ok"
-
     total_shifts = Rota.get_workers_total_shifts()
-
     for worker in Rota.get_workers():
         assert total_shifts[worker.name] == 4
         assert "aaaa" in Rota.get_worker_shift_list_string(worker)
 
-
+def test_worker_requirement2():
-    def test_worker_requirement2(self):
     Rota = generate_basic_rota(workers=2)
-
+    wr = [
-        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21)),
+        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 21)),
         WorkerRequirement(start_date=datetime.date(2022, 3, 28))
     ]
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                    assign_as_block=True, workers_required=wr),
-                assign_as_block=True,
-                #force_as_block=True,
-                workers_required=wr,
-                #start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
     )
     Rota.build_and_solve(options={"ratio": 0.1})
     Rota.export_rota_to_html("test_worker_requirement")
-
     assert Rota.results.solver.status == "ok"
-
     total_shifts = Rota.get_workers_total_shifts()
-
     for worker in Rota.get_workers():
         assert total_shifts[worker.name] == 18
         assert Rota.get_worker_shift_list_string(worker)[14:21] == "-------"
 
-    def test_worker_requirement_increase_workers(self):
+def test_worker_requirement_increase_workers():
     Rota = generate_basic_rota(workers=2)
-
+    wr = [
-        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
+        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
         WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
     ]
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                    assign_as_block=True, workers_required=wr),
-                assign_as_block=True,
-                #force_as_block=True,
-                workers_required=wr,
-                #start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
     )
     Rota.build_and_solve(options={"ratio": 0.1})
     Rota.export_rota_to_html("test_worker_requirement")
-
     assert Rota.results.solver.status == "ok"
-
     total_shifts = Rota.get_workers_total_shifts()
-
     for worker in Rota.get_workers():
         assert total_shifts[worker.name] == 30
         assert Rota.get_worker_shift_list_string(worker)[7:21] == "-" * 14
         assert Rota.get_worker_shift_list_string(worker).endswith("aaaa---"*7)
 
-    def test_worker_requirement_balancing(self):
+def test_worker_requirement_balancing():
     Rota = generate_basic_rota(workers=2)
-
     natws = generate_not_available_to_works(datetime.date(2022,4,18), days=21, reason="holiday")
-
+    Rota.add_worker(Worker(name="worker3", site="group1", grade=1, not_available_to_work=natws))
-        Rota.add_worker(
+    wr = [
-            Worker(name=f"worker3", site="group1", grade=1, not_available_to_work=natws)
+        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
-        )
-
-        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
         WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
     ]
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
+                    assign_as_block=True, workers_required=wr),
-                assign_as_block=True,
-                workers_required=wr,
-                #start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
     )
     Rota.build_and_solve(options={"ratio": 0.0})
     Rota.export_rota_to_html("test_worker_requirement")
-
     assert Rota.results.solver.status == "ok"
-
     total_shifts = Rota.get_workers_total_shifts()
-
     for worker in Rota.get_workers():
         assert total_shifts[worker.name] == 20
         assert Rota.get_worker_shift_list_string(worker)[7:21] == "-" * 14
-
         if worker.name == "worker3":
             assert Rota.get_worker_shift_list_string(worker).endswith("aaaa---")
             assert Rota.get_worker_shift_list_string(worker).startswith("aaaa---")
 
-
+def test_worker_requirement_multiple_shifts():
-    def test_worker_requirement_multiple_shifts(self):
     Rota = generate_basic_rota(workers=4)
-
+    wr = [
-        wr = [WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
+        WorkerRequirement(start_date=datetime.date(2022, 3, 7), end_date=datetime.date(2022, 3, 14)),
         WorkerRequirement(start_date=datetime.date(2022, 3, 28), number=2)
     ]
-
     Rota.add_shifts(
-            SingleShift(
+        SingleShift(sites=("group1", "group2"), name="b", length=12.5, days=days[4:],
-                sites=("group1", "group2"), name="b", length= 12.5, days=days[4:],
+                    assign_as_block=True, workers_required=2),
-                assign_as_block=True,
+        SingleShift(sites=("group1", "group2"), name="a", length=12.5, days=days[:4],
-                #force_as_block=True,
+                    assign_as_block=True, workers_required=wr),
-                workers_required=2,
-                #start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
-            SingleShift(
-                sites=("group1", "group2"), name="a", length= 12.5, days=days[:4],
-                assign_as_block=True,
-                #force_as_block=True,
-                workers_required=wr,
-                #start_date=(Rota.start_date + datetime.timedelta(days=7)),
-                #end_date=(Rota.start_date + datetime.timedelta(days=((7*6)-1))),
-            ),
     )
     Rota.build_and_solve(options={"ratio": 0.1})
     Rota.export_rota_to_html("test_worker_requirement")
-
     assert Rota.results.solver.status == "ok"
-
     total_shifts = Rota.get_workers_total_shifts()
-
     for worker in Rota.get_workers():
         assert total_shifts[worker.name] == 30
 
-class TestDisplayChar:
+def test_display_char_in_output(monkeypatch):
-    def test_display_char_in_output(monkeypatch):
     Rota = generate_basic_rota(workers=1)
     shift = SingleShift(
         sites=("group1",),
@@ -721,12 +448,8 @@ class TestDisplayChar:
         display_char="Z"
     )
     Rota.add_shifts(shift)
-
     assert shift.display_char == "Z"
-
     Rota.build_and_solve(options={"ratio": 0.1})
-
     assert Rota.results.solver.status == "ok"
-
     for worker in Rota.get_workers():
         assert Rota.get_worker_shift_list_string(worker).startswith("Z------" * 5)

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}
 
@@ -1024,9 +1052,9 @@ def test_worker_hard_day_dependency():
     worker2 = workers[1]
     worker3 = workers[2]
 
-    worker1.add_hard_day_dependency("Mon", "Tue")
+    worker1.add_hard_day_dependency({"Mon", "Tue"})
-    worker2.add_hard_day_dependency("Mon", "Wed")
+    worker2.add_hard_day_dependency({"Mon", "Wed"})
-    worker3.add_hard_day_dependency("Mon", "Thu")
+    worker3.add_hard_day_dependency({"Mon", "Thu"})
 
     Rota.add_shifts(
         SingleShift(
@@ -1080,15 +1108,15 @@ def test_worker_hard_day_dependency_weekend():
     worker3 = workers[2]
     worker4 = workers[3]
 
-    worker1.add_hard_day_dependency("Fri", "Sun")
+    worker1.add_hard_day_dependency({"Fri", "Sun"})
-    worker2.add_hard_day_dependency("Fri", "Sun")
+    worker2.add_hard_day_dependency({"Fri", "Sun"})
-    worker3.add_hard_day_dependency("Fri", "Sun")
+    worker3.add_hard_day_dependency({"Fri", "Sun"})
-    worker4.add_hard_day_dependency("Fri", "Sun")
+    worker4.add_hard_day_dependency({"Fri", "Sun"})
 
-    worker1.add_hard_day_exclusion("Fri", "Sat")
+    worker1.add_hard_day_exclusion({"Fri", "Sat"})
-    worker2.add_hard_day_exclusion("Fri", "Sat")
+    worker2.add_hard_day_exclusion({"Fri", "Sat"})
-    worker3.add_hard_day_exclusion("Fri", "Sat")
+    worker3.add_hard_day_exclusion({"Fri", "Sat"})
-    worker4.add_hard_day_exclusion("Fri", "Sat")
+    worker4.add_hard_day_exclusion({"Fri", "Sat"})
 
     Rota.add_shifts(
         SingleShift(
@@ -1128,15 +1156,15 @@ def test_worker_hard_day_dependency_weekend2():
     worker3 = workers[2]
     worker4 = workers[3]
 
-    worker1.add_hard_day_dependency("Fri", "Sun")
+    worker1.add_hard_day_dependency({"Fri", "Sun"})
-    worker2.add_hard_day_dependency("Fri", "Sun")
+    worker2.add_hard_day_dependency({"Fri", "Sun"})
-    worker3.add_hard_day_dependency("Fri", "Sun")
+    worker3.add_hard_day_dependency({"Fri", "Sun"})
-    worker4.add_hard_day_dependency("Fri", "Sun")
+    worker4.add_hard_day_dependency({"Fri", "Sun"})
-    worker4.add_hard_day_dependency("Fri", "Sat")
+    worker4.add_hard_day_dependency({"Fri", "Sat"})
 
-    worker1.add_hard_day_exclusion("Fri", "Sat")
+    worker1.add_hard_day_exclusion({"Fri", "Sat"})
-    worker2.add_hard_day_exclusion("Fri", "Sat")
+    worker2.add_hard_day_exclusion({"Fri", "Sat"})
-    worker3.add_hard_day_exclusion("Fri", "Sat")
+    worker3.add_hard_day_exclusion({"Fri", "Sat"})
 
     Rota.add_shifts(
         SingleShift(
@@ -1260,7 +1288,7 @@ def test_worker_force_assign_with_and_hard_days():
         worker.add_force_assign_with("a", "b")
 
         #worker.add_hard_day_dependency("Sat", "Sun")
-        worker.add_hard_day_dependency("Mon", "Tue")
+        worker.add_hard_day_dependency({"Mon", "Tue"})
 
     Rota.add_shifts(
         SingleShift(
@@ -1294,7 +1322,7 @@ def test_worker_force_assign_with_and_hard_days_weekend():
         worker.add_force_assign_with("a", "b")
 
         #worker.add_hard_day_dependency("Sat", "Sun")
-        worker.add_hard_day_dependency("Sat", "Sun")
+        worker.add_hard_day_dependency({"Sat", "Sun"})
 
     Rota.add_shifts(
         SingleShift(
@@ -1456,3 +1484,143 @@ def test_max_shifts_per_week_by_shift_name_invalid_shift():
     )
     with pytest.raises(InvalidShift):
         Rota.build_and_solve(options={"ratio": 0.0})
+
+def test_force_assign_shift():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=2)
+
+    workers = Rota.get_workers()
+
+    workers[0].force_assign_shift(1, "Mon", "a")
+    workers[0].force_assign_shift(1, "Tue", "b")
+    workers[0].force_assign_shift(2, "Wed", "b")
+    workers[0].force_assign_shift(2, "Thu", "b")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign", folder="tests")
+
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker)
+        if worker == workers[0]:
+            # Worker 0 should have 'a' on Monday and 'b' on Tuesday
+            assert shift_list[0] == "a", "Worker 0 should have 'a' on 1st Monday"
+            assert shift_list[1] == "b", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[9] == "b", "Worker 0 should have 'b' on 2nd Wednesday"
+            assert shift_list[10] == "b", "Worker 0 should have 'b' on 2nd Thursday"
+
+def test_force_assign_shift2():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=1)
+
+    workers = Rota.get_workers()
+
+    workers[0].force_assign_shift(1, "Mon", "a")
+    workers[0].force_assign_shift(1, "Tue", "a")
+    workers[0].force_assign_shift(1, "Wed", "a")
+    workers[0].force_assign_shift(1, "Thu", "a")
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+        SingleShift(
+            sites=("group1",),
+            name="b",
+            length=8,
+            days=days[:4],
+            workers_required=1,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_force_assign", folder="tests")
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker)
+        if worker == workers[0]:
+            # Worker 0 should have 'a' on Monday and 'b' on Tuesday
+            assert shift_list[0] == "a", "Worker 0 should have 'a' on 1st Monday"
+            assert shift_list[1] == "a", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[2] == "a", "Worker 0 should have 'b' on 1st Tuesday"
+            assert shift_list[3] == "a", "Worker 0 should have 'b' on 1st Tuesday"
+
+
+def test_worker_hard_day_exclusion():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_hard_day_exclusion({"Mon", "Tue"})
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[:5],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_hard_day_exclusion", folder="tests")
+
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker)
+        for week in range(16):
+            mon_idx = week * 7 + days.index("Mon")
+            tue_idx = week * 7 + days.index("Tue")
+            assert not (shift_list[mon_idx] == "a" and shift_list[tue_idx] == "a"), f"Worker {worker.name} has 'a' on both Mon and Tue in week {week}, which should not happen due to hard day exclusion"
+
+def test_worker_hard_day_exclusion2():
+    Rota = generate_basic_rota(workers=4, weeks_to_rota=16)
+
+    workers = Rota.get_workers()
+
+    for worker in workers:
+        worker.add_hard_day_exclusion({"Sat", "Sun"})
+
+        worker.add_hard_day_dependency({"Fri", "Sun"})
+
+    Rota.add_shifts(
+        SingleShift(
+            sites=("group1",),
+            name="a",
+            length=8,
+            days=days[4:],
+            workers_required=2,
+            assign_as_block=False,  # global setting is off
+        ),
+    )
+
+    Rota.build_and_solve(options={"ratio": 0.0})
+    Rota.export_rota_to_html("test_worker_hard_day_exclusion", folder="tests")
+
+    for worker in Rota.get_workers():
+        shift_list = Rota.get_worker_shift_list(worker)
+        for week in range(16):
+            sat_idx = week * 7 + days.index("Sat")
+            sun_idx = week * 7 + days.index("Sun")
+            assert not (shift_list[sat_idx] == "a" and shift_list[sun_idx] == "a"), f"Worker {worker.name} has 'a' on both Mon and Tue in week {week}, which should not happen due to hard day exclusion"