many improvements

gen.py

@@ -0,0 +1,311 @@
+import datetime
+import os
+import sys
+from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
+
+
+sites = (
+    "truro",
+    "exeter",
+    "torbay",
+    "barnstaple",
+    "plymouth",
+    "taunton",
+    "proc",
+    "truro no proc",
+)
+
+from rota.workers import (
+    Worker,
+    NotAvailableToWork,
+    NonWorkingDays,
+    WorkRequests,
+    PreferenceNotToWork,
+    OutOfProgramme,
+)
+
+
+suspend_on_finish = False
+solve = True
+time_to_run = 258500
+# allow = 5000
+ratio = 0.5
+
+
+start_date = datetime.date(2022, 9, 5)
+
+Rota = RotaBuilder(start_date, weeks_to_rota=12)
+Rota.constraint_options["balance_weekends"] = True
+Rota.constraint_options["max_night_frequency"] = 1
+Rota.constraint_options["max_weekend_frequency"] = 1
+# Rota.constraint_options["avoid_st2_first_month"] = True
+
+Rota.add_shifts(
+    SingleShift(
+        sites=("exeter", "exeter twilights and weekends"),
+        name="exeter_twilight",
+        length=12.5,
+        days=days[:5],
+        balance_offset=3,
+        constraint=[{"name": "max_shifts_per_week", "options": 2}],
+    ),
+    SingleShift(
+        sites=("truro", "truro twilights"),
+        name="truro_twilight",
+        length=12.5,
+        days=days[:4],
+        balance_offset=3,
+    ),
+    SingleShift(
+        sites=("torbay", "torbay twilights"),
+        name="torbay_twilight",
+        length=12.5,
+        days=days[:4],
+        balance_offset=5,
+        assign_as_block=True,
+    ),
+    SingleShift(
+        sites=("plymouth", "plymouth twilights"),
+        name="plymouth_twilight",
+        length=12.5,
+        days=days[:5],
+        balance_offset=3,
+        workers_required=1,
+        constraint=[{"name": "max_shifts_per_week", "options": 2}],
+    ),
+    SingleShift(
+        sites=("exeter", "exeter twilights and weekends"),
+        name="weekend_exeter",
+        length=12.5,
+        days=days[5:],
+        balance_offset=4,
+        rota_on_nwds=True,
+        force_as_block=True,
+        constraint=[{"name": "preclear2"}, {"name": "postclear2"}, {"name": "single_shift_block_per_week"}],
+        #constraint=[{"name": "preclear2"}, {"name": "postclear2"},],
+    ),
+    SingleShift(
+        sites=("truro",),
+        name="weekend_truro",
+        length=12.5,
+        days=days[4:],
+        balance_offset=4,
+        rota_on_nwds=True,
+        force_as_block=True,
+        constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+        # force_as_block_unless_nwd=True
+    ),
+    SingleShift(
+        sites=("torbay",),
+        name="weekend_torbay",
+        length=12.5,
+        days=days[4:],
+        balance_offset=4,
+        rota_on_nwds=True,
+        force_as_block=True,
+        constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+        # force_as_block_unless_nwd=True
+    ),
+    SingleShift(
+        sites=("plymouth",),
+        name="weekend_plymouth1",
+        length=8,
+        days=days[5:],
+        balance_offset=4,
+        workers_required=1,
+        rota_on_nwds=True,
+        force_as_block=True,
+        constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+    ),
+    SingleShift(
+        sites=("plymouth",),
+        name="weekend_plymouth2",
+        length=8,
+        days=days[5:],
+        balance_offset=4,
+        workers_required=1,
+        rota_on_nwds=True,
+        force_as_block=True,
+        constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
+    ),
+    SingleShift(
+        sites=("plymouth", "plymouth twilights"),
+        name="plymouth_bank_holidays",
+        length=8,
+        days=days[5:],
+        balance_offset=2,
+        workers_required=1,
+        rota_on_nwds=True,
+        force_as_block=False,
+        bank_holidays_only=True,
+    ),
+    SingleShift(
+        sites=(sites[:-1]),
+        name="night_weekday",
+        length=12.25,
+        days=days[:4],
+        balance_offset=4.9,
+        balance_weighting=1,
+        # hard_constrain_shift=False,
+        workers_required=3,
+        force_as_block=True,
+        rota_on_nwds=True,
+        constraint=[
+            {"name": "night"},
+            {"name": "preclear2"},
+            {"name": "postclear2"},
+            {"name": "require_remote_site_presence_week", "options": ("plymouth", 1)},
+        ],
+    ),
+    SingleShift(
+        sites=(sites[:-1]),
+        name="night_weekend",
+        length=12.25,
+        days=days[4:],
+        balance_offset=3.9,
+        balance_weighting=1,
+        # hard_constrain_shift=False,
+        workers_required=3,
+        force_as_block=True,
+        rota_on_nwds=True,
+        constraint=[
+            {"name": "night"},
+            {"name": "preclear2"},
+            {"name": "postclear2"},
+            {"name": "require_remote_site_presence_week", "options": ("plymouth", 1)},
+        ],
+    ),
+)
+
+Rota.add_grade_constraint_by_week([2], [1, 2, 3, 4], ["night_weekday", "night_weekend"])
+
+load_leave = True
+Rota.build_shifts()
+Rota.add_grade_constraint_by_week([2], [1, 2], Rota.get_shift_names())
+
+if load_leave:
+    import leave
+
+    workers = leave.load_leave(Rota)
+
+    n = 0
+    for worker in workers:
+        n = n + 1
+        worker_name = worker
+        w = workers[worker]
+        site = w["site"]
+        grade = w["grade"]
+        fte = float(w["fte"]) * 100
+        nwd = w["nwd"]
+        end_date = w["end_date"]
+        start_date = w["start_date"]
+        oop = w["oop"]
+        grade = w["grade"]
+        pair = w["pair"]
+        # name, site, grade, fte, nwd, end_date, oop, twi, twi_we, night, night_we, pw1, pw2, extra = row
+
+        # Ignore trainees if fte == 0 (what are they doing here anyway)
+        if fte == 0 or fte == "0" or fte == "":
+            continue
+
+        nwds = []
+        if nwd:
+            for i in nwd.split(","):
+                nwd_start_date = Rota.start_date
+                nwd_end_date = Rota.rota_end_date
+
+                if "[" in i:
+                    a, b = i.split("[")[1][:-1].split("-")
+                    nwd_start_date = datetime.datetime.strptime(a, "%d/%m/%y").date()
+                    nwd_end_date = datetime.datetime.strptime(b, "%d/%m/%y").date()
+
+                nwds.append(
+                    NonWorkingDays(
+                        day=i.split("[")[0].strip()[:3].capitalize(),
+                        start_date=nwd_start_date,
+                        end_date=nwd_end_date,
+                    )
+                )
+        # nwds = [i.strip()[:3].capitalize() for i in nwd.split("/")] if nwd else None
+        end_date = end_date if end_date else None
+        oop = oop.split("-") if oop else None
+
+        previous_shifts = {}
+        # if twi:
+        #    worked = twi.split(" ")[0]
+        #    allocated = twi.split(" ")[1][1:-1]
+        #    previous_shifts["{}_twilight".format(site.lower())] = (worked, allocated)
+        # if twi_we:
+        #    worked = twi_we.split(" ")[0]
+        #    allocated = twi_we.split(" ")[1][1:-1]
+        #    previous_shifts["weekend_{}".format(site.lower())] = (worked, allocated)
+        # if night:
+        #    worked = night.split(" ")[0]
+        #    allocated = night.split(" ")[1][1:-1]
+        #    previous_shifts["night_weekday".format(site.lower())] = (worked, allocated)
+        # if night_we:
+        #    worked = night_we.split(" ")[0]
+        #    allocated = night_we.split(" ")[1][1:-1]
+        #    previous_shifts["night_weekend".format(site.lower())] = (worked, allocated)
+        # if pw1:
+        #    worked = pw1.split(" ")[0]
+        #    allocated = pw1.split(" ")[1][1:-1]
+        #    previous_shifts["weekend_plymouth1"] = (worked, allocated)
+        # if pw2:
+        #    worked = pw2.split(" ")[0]
+        #    allocated = pw2.split(" ")[1][1:-1]
+        #    previous_shifts["weekend_plymouth2"] = (worked, allocated)
+
+        leave_requests = w["leave"]
+        work_requests = w["requests"]
+        site_pref = w["site_pref"]
+
+        Rota.add_worker(
+            Worker(
+                name=worker_name,
+                site=site.lower(),
+                grade=int(grade[2]),
+                id=n,
+                fte=int(fte),
+                nwds=nwds,
+                start_date=start_date,
+                end_date=end_date,
+                opp=oop,
+                not_available_to_work=leave_requests,
+                work_requests=work_requests,
+                remote_site=site_pref,
+                previous_shifts=previous_shifts,
+                pair=pair,
+                shift_balance_extra=w["shift_balance_extra"],
+                bank_holiday_extra=w["bank_holiday_extra"],
+            )
+        )
+
+Rota.build_workers()
+Rota.build_model()
+
+solver_options = {"ratio": ratio, "seconds": time_to_run, "threads": 10}
+
+Rota.solve_model(options=solver_options)
+
+# print(Rota.get_worker_details())
+
+# optimizer = SolverFactory('cbc')
+# result = optimizer.solve(prob,tee=True)
+# result.Solver.Status = SolverStatus.warning
+# prob.solutions.load_from(result)
+
+# ResultsHolder = RotaResults(Rota)
+
+# worker_timetable_brief = ResultsHolder.get_worker_timetable_brief(
+#    show_prefs=False, show_unavailable=False
+# )
+#
+# print(worker_timetable_brief)
+
+
+Rota.export_rota_to_html("proc_rota")
+# Rota.export_rota_to_csv("rota")
+
+if suspend_on_finish:
+    os.system("systemctl suspend")

leave.py

@@ -1,19 +1,25 @@
 import csv
+from datetime import datetime
 import re
+from requests import Session
+
+
+from rota.workers import Worker, NotAvailableToWork, NonWorkingDays, WorkRequests, PreferenceNotToWork, OutOfProgramme
 
 date_re = r"[\d]{1,2}\/[\d]{1,2}\/[\d]{2}"
 
 
-def load_leave(Rota, use_live_rota):
+def load_leave(Rota):
 
-    file_name = "leave.csv"
-
-    with open(file_name, newline="") as f:
+    with Session() as s:
         shifts = Rota.get_shift_names()
 
         workers = {}
 
-        reader = csv.reader(f, delimiter=",", quotechar='"')
+        download = s.get("https://docs.google.com/spreadsheets/d/e/2PACX-1vS49eVS0f86wUJIhXMq7tMPZfAEJZatER09JLu0xnlTjDBkC-54DkAPT3I0JVXDQXNG6AIh_fVeIRS3/pub?gid=0&single=true&output=csv")
+        decoded_content = download.content.decode('utf-8')
+
+        reader = csv.reader(decoded_content.splitlines(), delimiter=',')
 
         leave = {}
         requests = {}
@@ -21,7 +27,7 @@ def load_leave(Rota, use_live_rota):
 
         order = []
 
-        n = 0
+        n = 1
         for row in reader:
             row_title = row[0]
             r = row[2:]
@@ -71,8 +77,20 @@ def load_leave(Rota, use_live_rota):
                 elif row_title == "Flexible NWD":
                     worker["flexible_nwd"] = lower_item
 
-                elif row_title == "End Date":
-                    worker["end_date"] = lower_item
+                elif row_title in ("End Date", "CCT date"):
+                    if lower_item:
+                        date = datetime.strptime(lower_item, "%d/%m/%y").date()
+                        worker["end_date"] = date
+                    else:
+                        worker["end_date"] = None
+
+
+                elif row_title in ("Start date",):
+                    if lower_item:
+                        date = datetime.strptime(lower_item, "%d/%m/%y").date()
+                        worker["start_date"] = date
+                    else:
+                        worker["start_date"] = None
 
                 elif row_title == "OOP":
                     worker["oop"] = lower_item
@@ -97,18 +115,22 @@ def load_leave(Rota, use_live_rota):
 
                 elif row_title == "bank_holidays_worked":
                     if lower_item:
-                        print("BH", lower_item)
                         worker["bank_holiday_extra"] = int(lower_item)
 
                 elif re.match(date_re, row_title) is not None:
 
                     if lower_item != "":
+                        date = datetime.strptime(row[0], "%d/%m/%y").date()
                         # This may be easier to do as a dict
                         if lower_item in shifts:
-                            worker["requests"].append((row[0], lower_item))
+                            worker["requests"].append(WorkRequests(date=date, shift=lower_item))
+
+                        elif lower_item in ("work_request", "happy to work"):
+                            worker["requests"].append(WorkRequests(date=date, shift="*"))
+
 
                         else:
-                            worker["leave"].append((row[0], lower_item))
+                            worker["leave"].append(NotAvailableToWork(date=date, reason=lower_item))
 
             n = n + 1
 

output/timetable.css

@@ -18,6 +18,14 @@ table {
 
 }
 
+.shift-table {
+    padding-top: 10px;
+}
+
+.shift-table thead {
+    font-weight: bold;
+}
+
 .table-div {
     /* overflow-x: auto; */
     width: 80%;
@@ -107,10 +115,15 @@ th.bank-holiday {
     color: #0074D9;
 }
 
-.torquay {
+.torquay, .torbay {
     color: #5FA8E8;
 }
 
+.taunton {
+    color: #e85fdf;
+}
+
+
 .unavailable {
     color: lightgray;
 }
@@ -167,7 +180,7 @@ th.bank-holiday {
 }
 
 .nwd {
-    color: brown;
+    color: lightcoral;
 }
 
 .th {}
@@ -235,3 +248,54 @@ summary h2 {
     color: red;
 
 }
+
+/* table.transposed {
+    border-collapse: collapse;
+}
+
+table.transposed tr {
+    display: block;
+    float: left;
+}
+
+table.transposed th,
+table.transposed td {
+    display: block;
+    border: 1px solid black;
+    max-width: 100px;
+} */
+table.transposed {
+}
+
+table.transposed tr {
+}
+
+table.transposed th,
+table.transposed td {
+    width: fit-content;
+    max-width: unset;
+    position: relative;
+    border-top: solid 0.5px black;
+    border-right: solid 0.5px black;
+}
+
+table.transposed .Sat {
+    /* opacity: 30%; */
+    /* background-color: lightsteelblue; */
+    border-top: solid 2px lightsteelblue !important;
+}
+
+table.transposed .Sun {
+    border-bottom: solid 2px lightsteelblue !important;
+
+}
+
+table.transposed .bank-holiday {
+    border-bottom: dotted 2px blue !important;
+
+}
+
+table.transposed th.bank-holiday {
+    color: darkblue;
+
+}

output/timetable.js

@@ -161,7 +161,7 @@ function generateExtra() {
             <span>Total shifts: ${total_shifts}, </span>
             <span>Weekends: ${weekends_worked}, </span>
             <span>Bank holidays: ${bank_holidays_worked}, </span>
-            <span class='shift-diff-span'>Summed shift diff: ${summed_shift_diff}, </span>
+            <span class='shift-diff-span'>Summed shift diff: ${summed_shift_diff.toPrecision(2)}, </span>
             
             </div>`)
 
@@ -377,9 +377,61 @@ function viewWorker(worker) {
     </div>`)
     $("body").append(dlg)
 
+    console.log(ds.shiftCounts)
+    shifts = JSON.parse(ds.shiftCounts)
+    shift_targets = JSON.parse(ds.workerTargets)
+    shift_diffs = JSON.parse(ds.shiftDiff)
+
+    let shift_count_table = document.createElement("table")
+    shift_count_table.classList.add("shift-table")
+    let header_row = shift_count_table.createTHead().insertRow();
+    header_row.insertCell().appendChild(document.createTextNode("Shift"))
+    header_row.insertCell().appendChild(document.createTextNode("Count"))
+    header_row.insertCell().appendChild(document.createTextNode("Target"))
+    header_row.insertCell().appendChild(document.createTextNode("Diffs"))
+
+
+    table_body = shift_count_table.createTBody()
+
+    for (const shift in shifts) {
+        let row = table_body.insertRow();
+        row.insertCell().appendChild(document.createTextNode(shift))
+        row.insertCell().appendChild(document.createTextNode(shifts[shift]))
+        row.insertCell().appendChild(document.createTextNode(shift_targets[shift].toPrecision(2)))
+        row.insertCell().appendChild(document.createTextNode(shift_diffs[shift].toPrecision(2)))
+        //console.log(shift, shifts[shift]);
+    }
+
+    dlg.get(0).appendChild(shift_count_table)
     dlg.dialog()
 
 
     console.log(worker);
 
 }
+
+$("#export-div").append($("#main-table").clone().attr("id", "gen-table").addClass("transposed"))
+
+$("#gen-table").each(function() {
+        var $this = $(this);
+        var newrows = [];
+        //$this.find("th.site-title").parent().remove();
+        $this.find("tr").each(function(rowToColIndex){
+            $(this).find("td, th").each(function(colToRowIndex){
+                if(newrows[colToRowIndex] === undefined) { newrows[colToRowIndex] = $("<tr></tr>"); }
+                while(newrows[colToRowIndex].find("td, th").length < rowToColIndex){
+                    newrows[colToRowIndex].append($("<td></td>"));
+                }
+                newrows[colToRowIndex].append($(this));
+            });
+        });
+        $this.find("tr").remove();
+        $.each(newrows, function(){
+            $this.append(this);
+        });
+
+        $this.find("td.rota-day").each((n, el) => {
+            console.log(el)
+            $(el).text(el.dataset.shift)
+        });
+    });

requirements.txt

@@ -4,3 +4,9 @@ pytest
 black
 rich
 pydantic
+fastapi
+uvicorn
+fastapi-sqlalchemy
+alembic
+psycopg2
+python-dotenv

rota/run2.py

@@ -41,18 +41,12 @@ rota_collections["current"].add_shifts(
     SingleShift(("truro",), "weekend_truro", 12.5, days[5:], assign_as_block=True),
     SingleShift(("torquay",), "weekend_torquay", 12.5, days[5:], assign_as_block=True),
     SingleShift(
-        ("plymouth",),
-        "weekend_plymouth",
-        8,
-        days[5:],
+        sites=("plymouth",), name="weekend_plymouth", length= 8, days=days[5:],
         workers_required=2,
         assign_as_block=True,
     ),
     SingleShift(
-        (sites),
-        "night_weekday",
-        12.25,
-        days[:4],
+        sites=(sites), name="night_weekday", length= 12.25, days=days[:4],
         balance_offset=5,
         balance_weighting=1,
         workers_required=3,
@@ -61,10 +55,7 @@ rota_collections["current"].add_shifts(
         constraints=["night"],
     ),
     SingleShift(
-        (sites),
-        "night_weekend",
-        12.25,
-        days[4:],
+        sites=(sites), name="night_weekend", length= 12.25, days=days[4:],
         balance_offset=4,
         balance_weighting=1,
         workers_required=3,

rota/runfixnights.py

@@ -78,10 +78,7 @@ Rota.constraint_options["require_presence_at_site_overnight"] = ["plymouth"]
 
 Rota.add_shifts(
     SingleShift(
-        ("exeter",),
-        "exeter_twilight",
-        12.5,
-        days[:4],
+        sites=("exeter",), name="exeter_twilight", length= 12.5, days=days[:4],
         balance_offset=3,
         constraints=["max_2_shifts_per_week"],
     ),
@@ -89,76 +86,52 @@ Rota.add_shifts(
         ("truro", "truro no proc"), "truro_twilight", 12.5, days[:4], balance_offset=3
     ),
     SingleShift(
-        ("torbay", "torbay twilights"),
-        "torbay_twilight",
-        12.5,
-        days[:4],
+        sites=("torbay", "torbay twilights"), name="torbay_twilight", length= 12.5, days=days[:4],
         balance_offset=3,
         assign_as_block=True,
     ),
     SingleShift(
-        ("plymouth", "plymouth_twilights"),
-        "plymouth_twilight",
-        12.5,
-        days[:5],
+        sites=("plymouth", "plymouth_twilights"), name="plymouth_twilight", length= 12.5, days=days[:5],
         balance_offset=3,
         workers_required=1,
         constraints=["max_2_shifts_per_week"],
     ),
     SingleShift(
-        ("exeter",),
-        "weekend_exeter",
-        12.5,
-        days[4:],
+        sites=("exeter",), name="weekend_exeter", length= 12.5, days=days[4:],
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
     ),
     SingleShift(
-        ("truro", "truro no proc"),
-        "weekend_truro",
-        12.5,
-        days[4:],
+        sites=("truro", "truro no proc"), name="weekend_truro", length= 12.5, days=days[4:],
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
         # force_as_block_unless_nwd=True
     ),
     SingleShift(
-        ("torbay",),
-        "weekend_torbay",
-        12.5,
-        days[4:],
+        sites=("torbay",), name="weekend_torbay", length= 12.5, days=days[4:],
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
         # force_as_block_unless_nwd=True
     ),
     SingleShift(
-        ("plymouth",),
-        "weekend_plymouth1",
-        8,
-        days[5:],
+        sites=("plymouth",), name="weekend_plymouth1", length= 8, days=days[5:],
         balance_offset=4,
         workers_required=1,
         rota_on_nwds=True,
         force_as_block=True,
     ),
     SingleShift(
-        ("plymouth",),
-        "weekend_plymouth2",
-        8,
-        days[5:],
+        sites=("plymouth",), name="weekend_plymouth2", length= 8, days=days[5:],
         balance_offset=4,
         workers_required=1,
         rota_on_nwds=True,
         force_as_block=True,
     ),
     SingleShift(
-        ("plymouth", "plymouth_twilights"),
-        "plymouth_bank_holidays",
-        8,
-        days[5:],
+        sites=("plymouth", "plymouth_twilights"), name="plymouth_bank_holidays", length= 8, days=days[5:],
         balance_offset=2,
         workers_required=1,
         rota_on_nwds=True,
@@ -166,10 +139,7 @@ Rota.add_shifts(
         bank_holidays_only=True,
     ),
     SingleShift(
-        (sites[:-1]),
-        "night_weekday",
-        12.25,
-        days[:4],
+        sites=(sites[:-1]), name="night_weekday", length= 12.25, days=days[:4],
         balance_offset=4,
         balance_weighting=1,
         # hard_constrain_shift=False,
@@ -179,10 +149,7 @@ Rota.add_shifts(
         constraints=["night"],
     ),
     SingleShift(
-        (sites[:-1]),
-        "night_weekend",
-        12.25,
-        days[4:],
+        sites=(sites[:-1]), name="night_weekend", length= 12.25, days=days[4:],
         balance_offset=3,
         balance_weighting=1,
         # hard_constrain_shift=False,

rota/shifts.py

@@ -1,9 +1,9 @@
 from calendar import week
 import datetime
 import itertools
-from typing import Dict, Iterable, List, Sequence, Tuple, Set
+from typing import Dict, Iterable, List, Sequence, Tuple, Set, Any, Type
 
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, Extra, constr
 
 from pathlib import Path
 
@@ -47,6 +47,11 @@ SHIFT_BOUNDS = {
 }
 
 
+class ShiftConstraint(BaseModel):
+    name: str
+    options: bool | int | Dict | tuple = False
+
+
 class SingleShift(BaseModel):
     """Class to hold all details for a shift
 
@@ -70,40 +75,39 @@ class SingleShift(BaseModel):
 
 
     """
-    sites: Iterable[str]
+
+    sites: List[str]
     name: str
     length: float
-    shift_days: Sequence[str]
-    balance_offset: float | None = None,
-    balance_weighting: float = 1,
-    workers_required: float = 1,
-    rota_on_nwds: bool = False,
-    assign_as_block: bool = False,
-    force_as_block: bool = False,
-    force_as_block_unless_nwd: bool = False,
-    hard_constrain_shift: bool = True,
-    bank_holidays_only: bool = False,
-    constraint: Iterable[str] = [],
+    days: str | List[str]
+    balance_offset: float | None = None
+    balance_weighting: float = 1
+    workers_required: float = 1
+    rota_on_nwds: bool = False
+    assign_as_block: bool = False
+    force_as_block: bool = False
+    force_as_block_unless_nwd: bool = False
+    hard_constrain_shift: bool = True
+    bank_holidays_only: bool = False
+    constraint: List[ShiftConstraint] = []
 
     class Config:
         extra = Extra.allow
+        orm_mode = True
 
     def __init__(self, **data: Any):
         super().__init__(**data)
         # self.site = sites
         # self.name = name
         # self.length = length
-        #self.shift_days = shift_days
+        # self.days = days
         # self.balance_by_site = balance_by_site
         # balance_offset defines the max difference in allocated shifts
         # versus target shifts (hard constraint)
         # if 0 exactly equal shifts must be assigend (unlikely to be possible)
 
-
-
         if self.balance_offset is None:
-            self.balance_offset = len(self.shift_days)
-
+            self.balance_offset = len(self.days)
 
         if self.force_as_block_unless_nwd:
             self.assign_as_block = True
@@ -113,23 +117,33 @@ class SingleShift(BaseModel):
         self.constraint_options = {}
         self.constraints = []
         for c in self.constraint:
-            match c:
-                case (constraint, options):
-                    self.constraints.append(constraint)
-                    self.constraint_options[constraint] = options
-                case constraint:
-                    self.constraints.append(constraint)
+            self.constraints.append(c.name)
 
+            if c.options:
+                self.constraint_options[c.name] = c.options
+
+            # constraint = c.constraint
+        #            match c:
+        #                case (constraint, options):
+        #                    self.constraint_options[constraint] = options
+        #                case constraint:
+        #                    self.constraints.append(constraint)
+
+        if isinstance(self.days, str):
+            self.days = [self.days]
 
     def __str__(self):
-        return f"name: {self.name}, site: {self.site}, length: {self.length}, balance weighting: {self.balance_weighting}, balance offset: {self.balance_offset}, hard_constrain: {self.hard_constrain_shift}, workers required : {self.workers_required}, rota on nwds {self.rota_on_nwds}, assign as block: {self.assign_as_block}, force as block: {self.force_as_block}, constraints: {self.constraints}, constraint_options: {self.constraint_options}"  # , total shifts: {self.total_shifts}"
+        return f"name: {self.name}, site: {self.sites}, length: {self.length}, balance weighting: {self.balance_weighting}, balance offset: {self.balance_offset}, hard_constrain: {self.hard_constrain_shift}, workers required : {self.workers_required}, rota on nwds {self.rota_on_nwds}, assign as block: {self.assign_as_block}, force as block: {self.force_as_block}, constraints: {self.constraints}, constraint_options: {self.constraint_options}"  # , total shifts: {self.total_shifts}"
 
     def get_shift_summary(self):
         """
         Returns a text summary of the shift
         """
 
-        return f"{self.name}: {self.workers_required} workers for sites ({', '.join(self.site)}) on days ({', '.join(self.shift_days)})"
+        return f"{self.name}: {self.workers_required} workers for sites ({', '.join(self.site)}) on days ({', '.join(self.days)})"
+
+    def get_shift_number(self):
+        return len(self.days)
 
 
 class RotaBuilder(object):
@@ -137,16 +151,14 @@ class RotaBuilder(object):
 
     def __init__(
         self,
-        start_date: datetime.datetime = datetime.datetime.now(),
+        start_date: datetime.date = datetime.date.today(),
         weeks_to_rota: int = 26,
         balance_offset_modifier: int = 1,
         ltft_balance_offset: int = 1,
-        max_night_frequency: int = 2,
-        max_weekend_frequency: int = 1,  # Requires balance weekends
         use_previous_shifts: bool = False,
         use_shift_balance_extra: bool = False,
         use_bank_holiday_extra: bool = False,
-        SHIFT_BOUNDS = SHIFT_BOUNDS
+        SHIFT_BOUNDS=SHIFT_BOUNDS,
     ):
 
         print(f"Start time: {datetime.datetime.now()}")
@@ -155,7 +167,7 @@ class RotaBuilder(object):
         print(f"Use previous shifts (balance_extra) {use_shift_balance_extra}")
 
         self.SHIFT_BOUNDS = SHIFT_BOUNDS
-        self.shifts = []  # type List[SingleShift]
+        self.shifts: List[SingleShift] = []  # type List[SingleShift]
 
         self.days = days
 
@@ -180,7 +192,7 @@ class RotaBuilder(object):
         self.unavailable_to_work_reason = {}
         self.pref_not_to_work = {}
         self.pref_not_to_work_reason = {}
-        self.work_requests = set()
+        self.work_requests: set[Tuple[str, WeekInt, DayStr, ShiftName]] = set()
         self.work_requests_map = {}
 
         self.workers: list[Worker] = []
@@ -193,9 +205,6 @@ class RotaBuilder(object):
         self.balance_offset_modifier = balance_offset_modifier
         self.ltft_balance_offset = ltft_balance_offset
 
-        # Don't assign multiple shifts every (n) weeks
-        self.max_night_frequency = max_night_frequency
-        self.max_weekend_frequency = max_weekend_frequency
 
         self.constraint_options = {
             "balance_nights": True,
@@ -209,6 +218,10 @@ class RotaBuilder(object):
             "minimise_shift_diffs": False,  # less sophisticated version of balance_shifts_over_workers
             "balance_weekends": True,
             "max_weekends": 100,
+            # Don't assign multiple shifts every (n) weeks
+            "max_weekend_frequency": 1,  # Requires balance weekends
+            # Don't assign multiple shifts every (n) weeks
+            "max_night_frequency": 4,
             "max_shifts_per_week": 7,
             "max_shifts_per_month": 40,
             # The following will cause an unsolvable problem if a shift
@@ -219,6 +232,7 @@ class RotaBuilder(object):
             "prevent_thursdays_before_full_weekends": [],
             "avoid_st2_first_month": False,
             "hard_constrain_pair_separation": False,
+            "avoid_shifts_by_grades": [],
         }
 
         # Generate a map for week, day combinations to a given date
@@ -263,7 +277,9 @@ class RotaBuilder(object):
         if not results.solver.status:
             sys.exit(0)
 
-    def build_and_solve(self, options={"ratio": 0.1, "seconds": 1000, "threads": 10}, solve=True):
+    def build_and_solve(
+        self, options={"ratio": 0.1, "seconds": 1000, "threads": 10}, solve=True
+    ):
         self.build_shifts()
         self.build_workers()
         self.build_model()
@@ -611,11 +627,22 @@ class RotaBuilder(object):
             initialize=availability_init,
         )
 
+        work_request_sets = set()
+        for work_request in self.work_requests:
+            if work_request[3] == "*":
+                pass
+
+            else:
+                work_request_sets.add(work_request[3])
+
+
         # Validate shifts are valid
-        work_request_sets = set([i[3] for i in self.work_requests])
+        #work_request_sets = set([i[3] for i in self.work_requests])
         invalid_shifts = work_request_sets - set(self.get_shift_names())
         if invalid_shifts:
-            raise InvalidShift(f"Invalid worker shift request [shift(s): ${invalid_shifts}]")
+            raise InvalidShift(
+                f"Invalid worker shift request [shift(s): ${invalid_shifts}]"
+            )
 
         def work_request_init(model, wid, week, day, shift):
             if (wid, week, day, shift) in self.work_requests:
@@ -692,7 +719,6 @@ class RotaBuilder(object):
             #     self.model.works[worker.id, week, day, shift]
             #     for worker in self.workers if worker.site not in site_required))
 
-
         # Constraint: total hours worked hours worked
         # for worker in self.workers:
         #    self.model.constraints.add(
@@ -735,8 +761,7 @@ class RotaBuilder(object):
                 return Constraint.Skip
             return (
                 sum(
-                    model.shift_week_worker_assigned[shift.name, week, w.id]
-                    for w in workers
+                    model.shift_week_worker_assigned[shift, week, w.id] for w in workers
                 )
                 <= limit
             )
@@ -754,7 +779,7 @@ class RotaBuilder(object):
                     Constraint(
                         [week for week in self.weeks],
                         # [shift for shift in self.get_shifts_with_constraint("limit_grade_number")],
-                        [shift],
+                        [shift.name],
                         [grade],
                         [shift.constraint_options["limit_grade_number"][grade]],
                         rule=limitGradesRule,
@@ -769,8 +794,7 @@ class RotaBuilder(object):
                 raise ValueError("minimum_grade_number: not enough valid workers")
             return (
                 sum(
-                    model.shift_week_worker_assigned[shift.name, week, w.id]
-                    for w in workers
+                    model.shift_week_worker_assigned[shift, week, w.id] for w in workers
                 )
                 >= limit
             )
@@ -788,7 +812,7 @@ class RotaBuilder(object):
                 Constraint(
                     [week for week in self.weeks],
                     # [shift for shift in self.get_shifts_with_constraint("limit_grade_number")],
-                    [shift],
+                    [shift.name],
                     [grade],
                     [min_required],
                     rule=minimumGradesRule,
@@ -796,7 +820,9 @@ class RotaBuilder(object):
                 ),
             )
 
-        def presenceAtRemoteSiteWeek(model, week, shift, required_site, required_number):
+        def presenceAtRemoteSiteWeek(
+            model, week, shift, required_site, required_number
+        ):
             required_site_workers = [
                 w for w in self.workers if required_site == w.remote_site
             ]
@@ -809,7 +835,9 @@ class RotaBuilder(object):
                 >= required_number
             )
 
-        for shift in self.get_shifts_with_constraint("require_remote_site_presence_week"):
+        for shift in self.get_shifts_with_constraint(
+            "require_remote_site_presence_week"
+        ):
             site, required_number = shift.constraint_options[
                 "require_remote_site_presence_week"
             ]
@@ -826,7 +854,9 @@ class RotaBuilder(object):
                 ),
             )
 
-        def presenceAtRemoteSite(model, day, week, shift, required_site, required_number):
+        def presenceAtRemoteSite(
+            model, day, week, shift, required_site, required_number
+        ):
             required_site_workers = [
                 w for w in self.workers if required_site == w.remote_site
             ]
@@ -836,10 +866,7 @@ class RotaBuilder(object):
                     return Constraint.Skip
 
             return (
-                sum(
-                    model.works[w.id, week, day, shift]
-                    for w in required_site_workers
-                )
+                sum(model.works[w.id, week, day, shift] for w in required_site_workers)
                 >= required_number
             )
 
@@ -886,7 +913,7 @@ class RotaBuilder(object):
         #                )
         for site in self.sites:
             for shift in self.get_shifts_with_constraint("balance_across_groups"):
-                if site in shift.site:
+                if site in shift.sites:
                     block = shift.name
                     for week in self.weeks:
                         self.model.constraints.add(
@@ -919,7 +946,7 @@ class RotaBuilder(object):
                         ]
                         >= sum(
                             self.model.works[worker.id, week, day, shift_name]
-                            for day in shift.shift_days
+                            for day in shift.days
                         )
                     )
                 self.model.constraints.add(
@@ -941,8 +968,12 @@ class RotaBuilder(object):
                     for w in workers:
                         if w.non_working_day_list:
                             l = []
-                            for nwd, start_nwd_date, end_nwd_date in w.non_working_day_list:
-                                if nwd in shift.shift_days:
+                            for (
+                                nwd,
+                                start_nwd_date,
+                                end_nwd_date,
+                            ) in w.non_working_day_list:
+                                if nwd in shift.days:
                                     if start_nwd_date > self.get_week_start_date(week):
                                         continue
                                     if end_nwd_date < self.get_week_start_date(week):
@@ -957,8 +988,6 @@ class RotaBuilder(object):
                         else:
                             full_workers.append(w)
 
-                    print(nwd_workers)
-
                     if not nwd_workers:  # Just do the usual
                         self.model.constraints.add(
                             sum(
@@ -1015,6 +1044,33 @@ class RotaBuilder(object):
                     )
                 )
 
+            for constraint_dict in self.constraint_options["avoid_shifts_by_grades"]:
+                if invalid_shifts := set(constraint_dict["shifts"]).difference(
+                    set(self.get_shift_names())
+                ):
+                    raise ValueError(
+                        f"avoid shifts by grade constraint contains a non existent shift: {invalid_shifts}"
+                    )
+                if invalid_grades := set(constraint_dict["grades"]).difference(
+                    set(self.get_worker_grades())
+                ):
+                    raise ValueError(
+                        f"avoid shifts by grade constraint contains a non existent worker grade: {invalid_grades}"
+                    )
+
+                if worker.grade in constraint_dict["grades"]:
+                    self.model.constraints.add(
+                        0
+                        == sum(
+                            self.model.works[worker.id, week, day, shift]
+                            for week, day, shift in self.get_all_shiftname_combinations()
+                            if week in constraint_dict["weeks"]
+                            and shift in constraint_dict["shifts"]
+                        )
+                    )
+
+            # sys.exit(0)
+
             if self.constraint_options["avoid_st2_first_month"] and worker.grade == 2:
                 # Avoid ST1s on the first month
                 try:
@@ -1049,15 +1105,17 @@ class RotaBuilder(object):
             # no solution will be possible
             for shift in self.get_shifts():
                 if (
-                    worker.site in shift.site
+                    worker.site in shift.sites
                 ):  # Each site specfies which sites self.workers can fullfill it
 
                     # calaculate the total number of shifts that need assigning over the rota
-                    # total_shifts = (len(self.weeks) * len(shift.shift_days) *
-                    total_shifts = self.shift_counts[shift] * shift.workers_required
+                    # total_shifts = (len(self.weeks) * len(shift.days) *
+                    total_shifts = (
+                        self.shift_counts[shift.name] * shift.workers_required
+                    )
 
                     full_time_equivalent_joined = sum(
-                        self.full_time_equivalent_sites[i] for i in shift.site
+                        self.full_time_equivalent_sites[i] for i in shift.sites
                     )
 
                     target_shifts = (
@@ -1111,6 +1169,8 @@ class RotaBuilder(object):
                         )
                 else:
                     # Make sure shifts aren't assigned to those from other sites
+                    # Not needed if there are no shifts
+                    if self.get_week_day_combinations_for_shift(shift):
                         self.model.constraints.add(
                             0
                             == sum(
@@ -1147,7 +1207,7 @@ class RotaBuilder(object):
 
                 if self.constraint_options["balance_shifts_quadratic"]:
                     # This may need to be updated
-                    xU = 10
+                    xU = 25
                     xL = 1
                     self.model.constraints.add(
                         inequality(
@@ -1304,9 +1364,9 @@ class RotaBuilder(object):
             # We use a similar method to balance the number of weekends worked
             # This works as long as weekend shifts are assigned as blocks!
             weekend_shift_target_number = sum(
-                worker.shift_target_number[shift.name] / len(shift.shift_days)
+                worker.shift_target_number[shift.name] / len(shift.days)
                 for shift in self.get_shifts()
-                if "Sat" in shift.shift_days or "Sun" in shift.shift_days
+                if "Sat" in shift.days or "Sun" in shift.days
             )
 
             worker.weekend_shift_target_number = weekend_shift_target_number
@@ -1374,7 +1434,7 @@ class RotaBuilder(object):
 
             if self.constraint_options["balance_blocks"]:
                 for week_blocks in self.get_week_block_iterator(
-                    self.max_night_frequency
+                    self.constraint_options["max_night_frequency"]
                 ):
                     if self.get_shifts_with_constraint("night"):
                         # Prevent nights more than once every n weeks
@@ -1390,7 +1450,7 @@ class RotaBuilder(object):
                         )
 
             # if self.constraint_options["balance_weekends"]:
-            for week_blocks in self.get_week_block_iterator(self.max_weekend_frequency):
+            for week_blocks in self.get_week_block_iterator(self.constraint_options["max_weekend_frequency"]):
                 # Prevent weekend shifts more than once every n weeks
                 self.model.constraints.add(
                     1
@@ -1400,12 +1460,55 @@ class RotaBuilder(object):
                     )
                 )
 
+            for constraint_shift in self.get_shifts_with_constraints(
+                "max_shifts_per_week_block"
+            ):
+                constraint_options = constraint_shift.constraint_options[
+                    "max_shifts_per_week_block"
+                ]
+
+                # If not supplied, default to the number of days per shift
+                shift_number = constraint_shift.get_shift_number()
+                if "shift_number" in constraint_options:
+                    shift_number = constraint_options["shift_number"]
+
+                for week_blocks in self.get_week_block_iterator(
+                    constraint_options["weeks"]
+                ):
+                    # Prevent shifts more than once every n weeks
+                    self.model.constraints.add(
+                        shift_number
+                        >= sum(
+                            self.model.works[
+                                worker.id, week, day, constraint_shift.name
+                            ]
+                            for week in week_blocks
+                            for day in constraint_shift.days
+                        )
+                    )
+
             for week in self.weeks:
+                for constraint_shift in self.get_shifts_with_constraints(
+                    "max_shifts_per_week"
+                ):
+                    self.model.constraints.add(
+                        constraint_shift.constraint_options["max_shifts_per_week"]
+                        >= sum(
+                            self.model.works[worker.id, w, day, constraint_shift.name]
+                            # for shiftname in self.get_shift_names_by_week_day(week, day)
+                            # for day in self.days
+                            for w, day in self.get_week_day_combinations_for_shift(
+                                constraint_shift
+                            )
+                            if w == week
+                        )
+                    )
+
                 
                 self.model.constraints.add(
                     self.constraint_options["max_shifts_per_week"]
                     >= sum(
-                        self.model.works[worker.id, week, day, shiftname]
+                        self.model.works[worker.id, w, day, shiftname]
                         # for shiftname in self.get_shift_names_by_week_day(week, day)
                         # for day in self.days
                         for w, day, shiftname in self.get_all_shiftname_combinations(
@@ -1485,21 +1588,15 @@ class RotaBuilder(object):
                         full_weekend_count
                         >= sum(
                             self.model.works[worker.id, week, "Sat", shift.name]
-                            for shift in self.get_shift_names_by_week_day(
-                                week, "Sat"
-                            )
+                            for shift in self.get_shift_names_by_week_day(week, "Sat")
                         )
                         + sum(
                             self.model.works[worker.id, week, "Sun", shift.name]
-                            for shift in self.get_shift_names_by_week_day(
-                                week, "Sat"
-                            )
+                            for shift in self.get_shift_names_by_week_day(week, "Sat")
                         )
                         + sum(
                             self.model.works[worker.id, week, "Fri", shift.name]
-                            for shift in self.get_shift_names_by_week_day(
-                                week, "Sat"
-                            )
+                            for shift in self.get_shift_names_by_week_day(week, "Sat")
                         )
                     )
                 if (
@@ -1510,21 +1607,15 @@ class RotaBuilder(object):
                         full_weekend_count
                         >= sum(
                             self.model.works[worker.id, week, "Sat", shift.name]
-                            for shift in self.get_shift_names_by_week_day(
-                                week, "Sat"
-                            )
+                            for shift in self.get_shift_names_by_week_day(week, "Sat")
                         )
                         + sum(
                             self.model.works[worker.id, week, "Sun", shift.name]
-                            for shift in self.get_shift_names_by_week_day(
-                                week, "Sun"
-                            )
+                            for shift in self.get_shift_names_by_week_day(week, "Sun")
                         )
                         + sum(
                             self.model.works[worker.id, week, "Thu", shift.name]
-                            for shift in self.get_shift_names_by_week_day(
-                                week, "Thu"
-                            )
+                            for shift in self.get_shift_names_by_week_day(week, "Thu")
                         )
                     )
 
@@ -1576,6 +1667,27 @@ class RotaBuilder(object):
                 #     self.model.shift_week_worker_assigned[s.name, week, worker.id]
                 #     for s in self.get_shifts_with_constraint("night")))
 
+                #for shift in self.get_shifts_with_constraint("single_block_per_week"):
+                #    if shift.force_as_block:
+                #        # Force nights to be assigned in blocks
+                #        # self.model.constraints.add(8* self.model.shift_week_worker_assigned[shift.name, week, worker.id] >= sum(self.model.works[worker.id, week, day, shift.name] for day in self.days)
+                #        #     )
+
+                #        # Force night shifts to be assigned in blocks
+                #        self.model.constraints.add(
+                #            self.model.shift_week_worker_assigned[
+                #                shift.name, week, worker.id
+                #            ]
+                #            * len(shift.days)
+                #            == sum(
+                #                self.model.works[worker.id, w, d, s]
+                #                for w, d, s in self.get_all_shiftname_combinations(week=week)
+                #                #for day in self.days
+                #            )
+                #        )
+                #    else:
+                #        raise ValueError(f"Requires {shift.name} to have force_as_block")
+
                 # for shift in self.get_shifts_with_constraint("night"):
                 for shift in self.get_shifts():
                     if shift.force_as_block:
@@ -1588,10 +1700,10 @@ class RotaBuilder(object):
                             self.model.shift_week_worker_assigned[
                                 shift.name, week, worker.id
                             ]
-                            * len(shift.shift_days)
+                            * len(shift.days)
                             == sum(
                                 self.model.works[worker.id, week, day, shift.name]
-                                for day in shift.shift_days
+                                for day in shift.days
                             )
                         )
 
@@ -1603,6 +1715,9 @@ class RotaBuilder(object):
 
                 p1 = 1
                 p2 = 1
+                p3 = 1
+                p4 = 1
+                p5 = 1
                 if n > 0:
                     pweek, pday = weeks_days[n - 1]
                 else:
@@ -1615,6 +1730,24 @@ class RotaBuilder(object):
                     p2week, p2day = weeks_days[n]
                     p2 = 0
 
+                try:
+                    p3week, p3day = weeks_days[n - 3]
+                except IndexError:
+                    p3week, p3day = weeks_days[n]
+                    p3 = 0
+
+                try:
+                    p4week, p4day = weeks_days[n - 4]
+                except IndexError:
+                    p4week, p4day = weeks_days[n]
+                    p4 = 0
+
+                try:
+                    p5week, p5day = weeks_days[n - 5]
+                except IndexError:
+                    p5week, p5day = weeks_days[n]
+                    p5 = 0
+
                 n1 = 1
                 n2 = 1
                 try:
@@ -1679,7 +1812,7 @@ class RotaBuilder(object):
                 for constraint_shift in self.get_shifts_with_constraints(
                     "preclear", "preclear2"
                 ):
-                    if day in constraint_shift.shift_days:
+                    if day in constraint_shift.days:
                         self.model.constraints.add(
                             1
                             >= self.model.works[
@@ -1696,7 +1829,7 @@ class RotaBuilder(object):
                         )
 
                 for constraint_shift in self.get_shifts_with_constraint("preclear2"):
-                    if day in constraint_shift.shift_days:
+                    if day in constraint_shift.days:
                         self.model.constraints.add(
                             1
                             >= self.model.works[
@@ -1715,7 +1848,7 @@ class RotaBuilder(object):
                 for constraint_shift in self.get_shifts_with_constraints(
                     "postclear", "postclear2"
                 ):
-                    if day in constraint_shift.shift_days:
+                    if day in constraint_shift.days:
                         self.model.constraints.add(
                             1
                             >= self.model.works[
@@ -1732,7 +1865,7 @@ class RotaBuilder(object):
                         )
 
                 for constraint_shift in self.get_shifts_with_constraints("postclear2"):
-                    if day in constraint_shift.shift_days:
+                    if day in constraint_shift.days:
                         self.model.constraints.add(
                             1
                             >= self.model.works[
@@ -1967,7 +2100,6 @@ class RotaBuilder(object):
 
         Must be called prior to attempting to solve
         """
-        # self.build_shifts()
         if not self.workers:
             raise NoWorkers("Workers must be added prior to calling build_workers")
 
@@ -1982,11 +2114,11 @@ class RotaBuilder(object):
             self.workers_id_map[wid] = worker
 
             if worker.name in self.workers_name_map:
-                raise ValueError(f"Worker with name '{worker.name}' has been added twice")
+                raise ValueError(
+                    f"Worker with name '{worker.name}' has been added twice"
+                )
             self.workers_name_map[worker.name] = worker
 
-
-
         self.workers = sorted(self.workers)
 
         self.full_time_equivalent = sum(w.fte_adj for w in self.workers)
@@ -2007,7 +2139,14 @@ class RotaBuilder(object):
         for p in pairs:
             self.worker_pairs.append(tuple(pairs[p]))
 
-    def add_shift(self, shift) -> None:
+    def add_grade_constraint_by_week(
+        self, grades: Iterable[int], weeks: Iterable[int], shifts
+    ):
+        self.constraint_options["avoid_shifts_by_grades"].append(
+            {"grades": grades, "weeks": weeks, "shifts": shifts}
+        )
+
+    def add_shift(self, shift: SingleShift) -> None:
         """Add a shift to the collection
 
         :param SingleShift shift: Shift object
@@ -2023,7 +2162,7 @@ class RotaBuilder(object):
         """
         self.shifts.extend(shifts)
 
-    def get_shift_names_by_week_day(self, week, day: DayStr) -> set():
+    def get_shift_names_by_week_day(self, week, day: DayStr) -> set:
         """Returns the shifts required for a specific day
 
         Returns:
@@ -2066,6 +2205,9 @@ class RotaBuilder(object):
         self.week_day_shift_product = []
         self.week_day_shiftclass_product = []
 
+        for s in self.shifts:
+            pass
+
         for s in self.shifts:
             if s.name in self.shift_names:
                 raise InvalidShift(f"Duplicate shift: {s.name}")
@@ -2073,10 +2215,10 @@ class RotaBuilder(object):
             self.shifts_by_name[s.name] = s
             self.shift_names.append(s.name)
 
-            # for day in s.shift_days:
+            # for day in s.days:
             #    self.week_day_shifts_dict[(week, day)].add(s.name)
 
-            for site in s.site:
+            for site in list(s.sites):
                 self.sites.add(site)
 
         self.shift_counts = defaultdict(int)
@@ -2085,16 +2227,17 @@ class RotaBuilder(object):
             for s in self.shifts:
                 if s.bank_holidays_only:
                     if self.week_day_date_map[(week, day)] in bank_holiday_map:
+                        print(s.name, self.week_day_date_map[(week, day)])
                         self.week_day_shift_product.append((week, day, s.name))
                         self.week_day_shiftclass_product.append((week, day, s))
                         self.week_day_shifts_dict[(week, day)].add(s.name)
-                        self.shift_counts[s] = self.shift_counts[s] + 1
+                        self.shift_counts[s.name] = self.shift_counts[s.name] + 1
                 else:
-                    if day in s.shift_days:
+                    if day in s.days:
                         self.week_day_shift_product.append((week, day, s.name))
                         self.week_day_shiftclass_product.append((week, day, s))
                         self.week_day_shifts_dict[(week, day)].add(s.name)
-                        self.shift_counts[s] = self.shift_counts[s] + 1
+                        self.shift_counts[s.name] = self.shift_counts[s.name] + 1
         # print(self.week_day_shift_product)
 
     #        # Todo replace with week_day.....
@@ -2102,7 +2245,7 @@ class RotaBuilder(object):
     #        self.day_shiftclass_product = []
     #        for s in self.shifts:
     #            for day in days:
-    #                if day in s.shift_days:
+    #                if day in s.days:
     #                    self.day_shift_product.append((day, s.name))
     #                    self.day_shiftclass_product.append((day, s))
     #
@@ -2193,17 +2336,27 @@ class RotaBuilder(object):
         """
         return self.shifts
 
+    def get_shifts_for_worker_site(self, worker_site):
+        shifts = [shift for shift in self.shifts if worker_site in shift.sites]
+        return shifts
+
+    def get_shifts_for_worker(self, worker_id):
+        worker = self.get_worker_by_id(worker_id)
+        shifts = [shift for shift in self.shifts if worker.site in shift.sites]
+        return shifts
+
     def get_shifts_with_constraint(self, constraint) -> List[SingleShift]:
         return [shift for shift in self.shifts if constraint in shift.constraints]
 
-    def get_shifts_with_constraints(self, *constraints) -> set[SingleShift]:
-        shifts = set()
+    def get_shifts_with_constraints(self, *constraints) -> List[SingleShift]:
+        shift_names = set()
 
         for constraint in constraints:
-            shifts.update(
-                [shift for shift in self.shifts if constraint in shift.constraints]
+            shift_names.update(
+                [shift.name for shift in self.shifts if constraint in shift.constraints]
             )
-        return shifts
+
+        return [self.get_shift_by_name(s) for s in shift_names]
 
     def get_shift_names(self) -> List[ShiftName]:
         """Returns a list of all the registered shift names
@@ -2239,12 +2392,12 @@ class RotaBuilder(object):
         """Returns a list of all possible shifts, the workers required and site
 
         Returns:
-            list: list (week, day, shift name, workers required, shift site)
+            list: list (week, day, shift name, workers required, shift.sites)
         """
         l = []
         for week, day, shift in self.week_day_shiftclass_product:
-            # if day in shift.shift_days:
-            l.append((week, day, shift.name, shift.workers_required, shift.site))
+            # if day in shift.days:
+            l.append((week, day, shift.name, shift.workers_required, shift.sites))
 
         return l
 
@@ -2265,17 +2418,24 @@ class RotaBuilder(object):
 
         return l - set(self.get_all_shiftclass_combinations())
 
-    def shifts_to_assign_as_blocks(self):
+    def shifts_to_assign_as_blocks(self) -> List[str]:
         return [shift.name for shift in self.get_shifts() if shift.assign_as_block]
 
-    def shifts_to_force_as_blocks(self):
-        return [shift.name for shift in self.get_shifts() if (shift.force_as_block or shift.force_as_block_unless_nwd)]
+    def shifts_to_force_as_blocks(self) -> List[str]:
+        return [
+            shift.name
+            for shift in self.get_shifts()
+            if (shift.force_as_block or shift.force_as_block_unless_nwd)
+        ]
 
-    def shifts_to_assign_or_force_as_blocks(self):
+    def shifts_to_assign_or_force_as_blocks(self) -> List[str]:
         s = self.shifts_to_assign_as_blocks()
         s.extend(self.shifts_to_force_as_blocks())
         return s
 
+    def get_worker_by_id(self, worker_id: str) -> Worker:
+        return self.workers_id_map[worker_id]
+
     def get_worker_by_name(self, name: str) -> Worker:
         return self.workers_name_map[name]
 
@@ -2301,7 +2461,7 @@ class RotaBuilder(object):
         return group_workers
 
     def get_workers_for_shift(self, shift: SingleShift) -> List[Worker]:
-        return [worker for worker in self.workers if worker.site in shift.site]
+        return [worker for worker in self.workers if worker.site in shift.sites]
 
     def get_workers_total_fte(self) -> float:
         return self.full_time_equivalent
@@ -2322,6 +2482,9 @@ class RotaBuilder(object):
     def get_workers(self):
         return self.workers
 
+    def get_worker_grades(self) -> set[int]:
+        return set([worker.grade for worker in self.workers])
+
     def get_bank_holiday_week_days(self):
         return [
             (week, day)
@@ -2385,7 +2548,6 @@ class RotaBuilder(object):
                         week_table[week][day][shift].append(worker.get_details())
         return week_table
 
-
     def get_worker_timetable(self):
         works = self.model.works
         timetable = {
@@ -2481,7 +2643,9 @@ class RotaBuilder(object):
             if d in bank_holiday_map:
                 th_class = "bank-holiday"
 
-            date_row.append(f"<th title='{d}' class='{th_class}' data-date='{d}'>Week {week}: {day}</th>")
+            date_row.append(
+                f"<th title='{d}' class='{th_class}' data-date='{d}'>Week {week}: {day}</th>"
+            )
             n = n + 1
         timetable.append(f"<tr class='data-row'>{''.join(date_row)}</tr>")
 
@@ -2667,7 +2831,7 @@ class RotaBuilder(object):
         html = f""" 
         <body>
         <div class="table-div" id="{table_name}">
-        <table>{joined_timetable}</table>
+        <table id="main-table">{joined_timetable}</table>
         </div>
         <details>
         <summary><h2>Rota settings</h2></summary>
@@ -2688,7 +2852,14 @@ class RotaBuilder(object):
         <pre>
         {result_string}
         </pre>
+        </details>
         <div>
+        <details>
+        <summary><h2>Export table</h2></summary>
+        <div id="export-div">
+        <div>
+        </details>
+        </div
         </body>
         """
 
@@ -2738,12 +2909,12 @@ class RotaBuilder(object):
         shifts = {}
 
         for worker in self.workers:
-            shifts[worker.name] = len([i for i in self.get_worker_shift_list(worker) if i != ""])
+            shifts[worker.name] = len(
+                [i for i in self.get_worker_shift_list(worker) if i != ""]
+            )
 
         return shifts
 
-
-
     def get_worker_shift_list(self, worker: Worker) -> List:
         shifts = []
 
@@ -2822,14 +2993,17 @@ class RotaBuilder(object):
 
 class NoActiveSites(Exception):
     """Raised when there are no active sites"""
+
     pass
 
 
 class NoWorkers(Exception):
     """Raised when there are no active sites"""
+
     pass
 
 
 class InvalidShift(Exception):
     """Raised when there are no active sites"""
+
     pass

rota/workers.py

@@ -1,7 +1,7 @@
 import datetime
 from collections import defaultdict
 from typing import Iterable, List, Literal, Optional
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel, Extra, validator
 
 from rich.pretty import pprint
 
@@ -15,29 +15,51 @@ if TYPE_CHECKING:
     from rota.shifts import RotaBuilder
 
 days = Literal["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
+whole_days = [
+    "sunday",
+    "monday",
+    "tuesday",
+    "wednesday",
+    "thursday",
+    "friday",
+    "saturday",
+]
+
 
 class NotAvailableToWork(BaseModel):
     date: datetime.date
     reason: str = "unknown"
 
+
 class NonWorkingDays(BaseModel):
     day: days
     start_date: datetime.date | None = None
     end_date: datetime.date | None = None
 
+    @validator("day")
+    def day_in(cls, day):
+        for whole_day in whole_days:
+            if day.lower() in whole_day:
+                return whole_day[:3].capitalize()
+        raise ValueError("Invalid day")
+
+
 class WorkRequests(BaseModel):
     date: datetime.date
     shift: str = ""
 
+
 class PreferenceNotToWork(BaseModel):
     date: datetime.date
     reason: str = ""
 
+
 class OutOfProgramme(BaseModel):
     start_date: datetime.date
     end_date: datetime.date
     reason: str = ""
 
+
 class Worker(BaseModel):
     name: str
     site: str
@@ -54,7 +76,7 @@ class Worker(BaseModel):
     not_available_to_work: list[NotAvailableToWork] = []
     pref_not_to_work: list[PreferenceNotToWork] = []
     work_requests: list[WorkRequests] = []
-    remote_site: str = "plymouth",  # We set a default proc_site
+    remote_site: str = ("plymouth",)  # We set a default proc_site
     previous_shifts: dict = {}
     shift_balance_extra: dict = {}
     bank_holiday_extra: int = 0
@@ -96,6 +118,9 @@ class Worker(BaseModel):
             if self.start_date >= self.end_date:
                 raise ValueError("End date must be after start date")
 
+        if self.name == "Steph bailey":
+            print(self.start_date, self.end_date)
+
         # if no start date default to the start of the rota
         if self.start_date is None:
             self.calculated_start_date = Rota.start_date
@@ -144,9 +169,7 @@ class Worker(BaseModel):
             if isinstance(end_oop, datetime.date):
                 end_oop_date = end_oop
             else:
-                end_oop_date = datetime.datetime.strptime(
-                    end_oop, "%d/%m/%y"
-                ).date()
+                end_oop_date = datetime.datetime.strptime(end_oop, "%d/%m/%y").date()
 
             if start_oop_date >= end_oop_date:
                 raise ValueError("End OOP date must be after start date")
@@ -178,10 +201,7 @@ class Worker(BaseModel):
 
         # loop throught dates converting to week / day combination
         for item in self.pref_not_to_work:
-            days_from_start = (
-                item.date
-                - Rota.start_date
-            ).days
+            days_from_start = (item.date - Rota.start_date).days
             week = days_from_start // 7 + 1
             day = Rota.days[(days_from_start % 7)]
             # Weight the value to take into account the number of preferences
@@ -195,22 +215,20 @@ class Worker(BaseModel):
         # print(not_available_to_work)
         # loop throught dates converting to week / day combination
         for item in self.not_available_to_work:
-            days_from_start = (
-                item.date
-                - Rota.start_date
-            ).days
+            days_from_start = (item.date - Rota.start_date).days
             week = days_from_start // 7 + 1
             day = Rota.days[(days_from_start % 7)]
             Rota.unavailable_to_work.add((self.id, week, day))
             Rota.unavailable_to_work_reason[(self.id, week, day)] = item.reason
 
         for item in self.work_requests:
-                days_from_start = (
-                    item.date
-                    - Rota.start_date
-                ).days
+            days_from_start = (item.date - Rota.start_date).days
             week = days_from_start // 7 + 1
             day = Rota.days[(days_from_start % 7)]
+            if item.shift == "*":
+                for shift in Rota.get_shifts_for_worker_site(self.site):
+                    Rota.work_requests.add((self.id, week, day, shift.name))
+            else:
                 Rota.work_requests.add((self.id, week, day, item.shift))
 
         # Calculate the proportion of the rota that is being worked
@@ -236,6 +254,9 @@ class Worker(BaseModel):
 
             self.non_working_day_list.append((item.day, start_date, end_date))
 
+        if days_to_work < 1:
+            self.fte_adj = 0
+
     def __lt__(self, other) -> bool:
         return (self.site, self.grade, self.fte_adj, self.name) < (
             other.site,
@@ -246,7 +267,11 @@ class Worker(BaseModel):
 
     def __str__(self) -> str:
 
-        nwds = ", ".join([str(i) for i, start, end in self.non_working_day_list]) if self.non_working_day_list is not None else ""
+        nwds = (
+            ", ".join([str(i) for i, start, end in self.non_working_day_list])
+            if self.non_working_day_list is not None
+            else ""
+        )
         return "{} {} [{}] REMOTE SITE:{}".format(
             self.name,
             (self.site, self.grade, self.fte_adj),

run.py

@@ -78,10 +78,7 @@ Rota.constraint_options["max_weekends"] = 7
 
 Rota.add_shifts(
     SingleShift(
-        ("exeter",),
-        "exeter_twilight",
-        12.5,
-        days[:4],
+        sites=("exeter",), name="exeter_twilight", length= 12.5, days=days[:4],
         balance_offset=3,
         constraints=["max_2_shifts_per_week"],
     ),
@@ -89,37 +86,25 @@ Rota.add_shifts(
         ("truro", "truro no proc"), "truro_twilight", 12.5, days[:4], balance_offset=3
     ),
     SingleShift(
-        ("torbay", "torbay twilights"),
-        "torbay_twilight",
-        12.5,
-        days[:4],
+        sites=("torbay", "torbay twilights"), name="torbay_twilight", length= 12.5, days=days[:4],
         balance_offset=5,
         assign_as_block=True,
     ),
     SingleShift(
-        ("plymouth", "plymouth_twilights"),
-        "plymouth_twilight",
-        12.5,
-        days[:5],
+        sites=("plymouth", "plymouth_twilights"), name="plymouth_twilight", length= 12.5, days=days[:5],
         balance_offset=3,
         workers_required=1,
         constraints=["max_2_shifts_per_week"],
     ),
     SingleShift(
-        ("exeter",),
-        "weekend_exeter",
-        12.5,
-        days[4:],
+        sites=("exeter",), name="weekend_exeter", length= 12.5, days=days[4:],
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
         constraints=["preclear2", "postclear2"],
     ),
     SingleShift(
-        ("truro", "truro no proc"),
-        "weekend_truro",
-        12.5,
-        days[4:],
+        sites=("truro", "truro no proc"), name="weekend_truro", length= 12.5, days=days[4:],
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
@@ -127,10 +112,7 @@ Rota.add_shifts(
         # force_as_block_unless_nwd=True
     ),
     SingleShift(
-        ("torbay",),
-        "weekend_torbay",
-        12.5,
-        days[4:],
+        sites=("torbay",), name="weekend_torbay", length= 12.5, days=days[4:],
         balance_offset=4,
         rota_on_nwds=True,
         force_as_block=True,
@@ -138,10 +120,7 @@ Rota.add_shifts(
         # force_as_block_unless_nwd=True
     ),
     SingleShift(
-        ("plymouth",),
-        "weekend_plymouth1",
-        8,
-        days[5:],
+        sites=("plymouth",), name="weekend_plymouth1", length= 8, days=days[5:],
         balance_offset=4,
         workers_required=1,
         rota_on_nwds=True,
@@ -149,10 +128,7 @@ Rota.add_shifts(
         constraints=["preclear2", "postclear2"],
     ),
     SingleShift(
-        ("plymouth",),
-        "weekend_plymouth2",
-        8,
-        days[5:],
+        sites=("plymouth",), name="weekend_plymouth2", length= 8, days=days[5:],
         balance_offset=4,
         workers_required=1,
         rota_on_nwds=True,
@@ -160,10 +136,7 @@ Rota.add_shifts(
         constraints=["preclear2", "postclear2"],
     ),
     SingleShift(
-        ("plymouth", "plymouth_twilights"),
-        "plymouth_bank_holidays",
-        8,
-        days[5:],
+        sites=("plymouth", "plymouth_twilights"), name="plymouth_bank_holidays", length= 8, days=days[5:],
         balance_offset=2,
         workers_required=1,
         rota_on_nwds=True,
@@ -171,10 +144,7 @@ Rota.add_shifts(
         bank_holidays_only=True,
     ),
     SingleShift(
-        (sites[:-1]),
-        "night_weekday",
-        12.25,
-        days[:4],
+        sites=(sites[:-1]), name="night_weekday", length= 12.25, days=days[:4],
         balance_offset=4.9,
         balance_weighting=1,
         # hard_constrain_shift=False,
@@ -184,10 +154,7 @@ Rota.add_shifts(
         constraints=["night", "preclear2", "postclear2"],
     ),
     SingleShift(
-        (sites[:-1]),
-        "night_weekend",
-        12.25,
-        days[4:],
+        sites=(sites[:-1]), name="night_weekend", length= 12.25, days=days[4:],
         balance_offset=3.9,
         balance_weighting=1,
         # hard_constrain_shift=False,

test/rota_constraints/test_avoid_grades_by_week.py

@@ -0,0 +1,195 @@
+import datetime
+import pytest
+from rota.shifts import InvalidShift, NoWorkers, RotaBuilder, SingleShift, days
+
+from rota.workers import Worker
+
+import itertools
+
+
+def generate_basic_rota(weeks_to_rota=9):
+    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="worker1", site="group1", grade=1),
+            Worker(name="worker2", site="group1", grade=2),
+            Worker(name="worker3", site="group1", grade=3),
+        ]
+    )
+
+    return Rota
+
+
+def date_generator(from_date, days):
+    n = 0
+    while True:
+        yield from_date
+
+        n = n + 1
+        if n >= days:
+            break
+        from_date = from_date + datetime.timedelta(days=1)
+
+
+class TestWorkerRequests:
+    def test_avoid_grade_constraint(self):
+        Rota = generate_basic_rota()
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=2,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.add_grade_constraint_by_week([1], [1,2,3], ["a"])
+        Rota.add_grade_constraint_by_week([2], [4,5,6], ["a"])
+        Rota.add_grade_constraint_by_week([3], [7,8,9], ["a"])
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("avoid_grade_constraint")
+
+        shift_summary = Rota.get_shift_summary_dict()
+        for worker_name in shift_summary:
+            worker_shifts = shift_summary[worker_name]
+
+            worker = Rota.get_worker_by_name(worker_name)
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            match worker_name:
+                case "worker1":
+                    assert shift_string.startswith("-"*21)
+                case "worker2":
+                    assert shift_string.startswith(21*"a"+"-"*21)
+                case "worker3":
+                    assert shift_string.endswith("-"*21)
+
+
+    def test_avoid_grade_constraint_invalid_shift(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+        )
+        with pytest.raises(ValueError):
+
+            Rota.add_grade_constraint_by_week([1], [1,2,3], ["b"])
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="b", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.add_grade_constraint_by_week([1], [1,2,3], ["b"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "ok"
+
+
+    def test_avoid_grade_constraint_invalid_grade(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=2,
+                force_as_block=False,
+            ),
+        )
+        with pytest.raises(ValueError):
+
+            Rota.add_grade_constraint_by_week([4], [1,2,3], ["a"])
+            Rota.build_and_solve(options={"ratio": 0.000})
+
+        Rota.add_worker(
+                Worker(name="worker4", site="group1", grade=4),
+        )
+
+
+        Rota.add_grade_constraint_by_week([4], [1,2,3], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "ok"
+
+    def test_avoid_grade_constraint_multiple_grades(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.add_grade_constraint_by_week([1,2], [1,2,3], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "ok"
+
+        Rota.add_grade_constraint_by_week([3], [1,2,3], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "warning"
+
+    def test_avoid_grade_constraint_multiple_shifts(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1",), name="a", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+            SingleShift(
+                sites=("group1",), name="b", length= 12.5, days=days,
+                workers_required=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.add_grade_constraint_by_week([1], [1,2,3], ["a", "b"])
+        Rota.add_grade_constraint_by_week([2], [4,5,6], ["a", "b"])
+        Rota.add_grade_constraint_by_week([3], [7,8,9], ["a", "b"])
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("avoid_grade_constraint")
+
+        shift_summary = Rota.get_shift_summary_dict()
+        for worker_name in shift_summary:
+            worker_shifts = shift_summary[worker_name]
+
+            worker = Rota.get_worker_by_name(worker_name)
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            match worker_name:
+                case "worker1":
+                    assert shift_string.startswith("-"*21)
+                case "worker2":
+                    assert "-"*21 in shift_string # hacky
+                case "worker3":
+                    assert shift_string.endswith("-"*21)
+
+        Rota.add_grade_constraint_by_week([2], [7,8,9], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "ok"
+
+        Rota.add_grade_constraint_by_week([1], [7,8,9], ["a"])
+        Rota.build_and_solve(options={"ratio": 0.000})
+        assert Rota.results.solver.status == "warning"

test/test_general_balancing.py

@@ -15,14 +15,16 @@ def generate_basic_rota(weeks_to_rota=10):
     Rota.constraint_options["balance_weekends"] = True
 
     # Add a few workers
-    Rota.add_workers([
+    Rota.add_workers(
+        [
             Worker(name="worker1", site="group1", grade=1),
             Worker(name="worker2", site="group1", grade=1),
             Worker(name="worker3", site="group2", grade=1),
             # Worker(name="worker4", site="group2", grade=1),
             # Worker(name="worker5", site="group2", grade=1),
             # Worker(name="worker6", site="group2", grade=1, fte=50),
-    ])
+        ]
+    )
 
     return Rota
 
@@ -33,22 +35,22 @@ class TestBalancing:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days,
                 balance_offset=2,
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
             SingleShift(
-                ("group1", "group2"),
-                "b",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=days,
                 balance_offset=2,
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
         )
 
@@ -68,30 +70,30 @@ class TestBalancing:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=["group1", "group2"],
+                name="a",
+                length=12.5,
+                days=days,
                 balance_weighting=4,
                 workers_required=1,
                 force_as_block=False,
-                constraint=["preclear2", "postclear2"],
+                constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
             ),
             SingleShift(
-                ("group1", "group2"),
-                "b",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=days,
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
             SingleShift(
-                ("group1", "group2"),
-                "c",
-                12.5,
-                days[0],
+                sites=("group1", "group2"),
+                name="c",
+                length=12.5,
+                days=days[0],
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
         )
 
@@ -111,31 +113,31 @@ class TestBalancing:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days,
                 # balance_weighting=4,
                 workers_required=1,
                 force_as_block=False,
                 # constraint=["preclear2", "postclear2"],
             ),
             SingleShift(
-                ("group1", "group2"),
-                "b",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=days,
                 balance_weighting=4,
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
             SingleShift(
-                ("group1", "group2"),
-                "c",
-                12.5,
-                days[0],
+                sites=("group1", "group2"),
+                name="c",
+                length=12.5,
+                days=days[0],
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
         )
 
@@ -156,32 +158,32 @@ class TestBalancing:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days,
                 # balance_weighting=4,
                 workers_required=1,
                 force_as_block=False,
                 # constraint=["preclear2", "postclear2"],
             ),
             SingleShift(
-                ("group1", "group2"),
-                "b",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=days,
                 balance_weighting=4,
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
             SingleShift(
-                ("group1", "group2"),
-                "c",
-                12.5,
-                days[0],
+                sites=("group1", "group2"),
+                name="c",
+                length=12.5,
+                days=days[0],
                 balance_weighting=4,
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
         )
 
@@ -206,10 +208,10 @@ class TestBalancing:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days,
                 balance_offset=99,
                 balance_weighting=8,
                 workers_required=2,
@@ -217,14 +219,14 @@ class TestBalancing:
                 # constraint=["preclear2", "postclear2"],
             ),
             SingleShift(
-                ("group2","group3"),
-                "b",
-                12.5,
-                days,
+                sites=("group2", "group3"),
+                name="b",
+                length=12.5,
+                days=days,
                 balance_offset=99,
                 # balance_weighting=4,
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
         )
 
@@ -247,10 +249,10 @@ class TestBalancing:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days,
                 balance_offset=99,
                 # balance_weighting=8,
                 workers_required=2,
@@ -258,14 +260,14 @@ class TestBalancing:
                 # constraint=["preclear2", "postclear2"],
             ),
             SingleShift(
-                ("group2","group3"),
-                "b",
-                12.5,
-                days,
+                sites=("group2", "group3"),
+                name="b",
+                length=12.5,
+                days=days,
                 balance_offset=99,
                 balance_weighting=4,
                 workers_required=1,
-                force_as_block=False
+                force_as_block=False,
             ),
         )
 

test/test_nwds.py

@@ -32,25 +32,19 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 # balance_offset=10,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "w",
-                12.5,
-                days[5:],
+                sites=("group1",), name="w", length= 12.5, days=days[5:],
                 workers_required=2,
             ),
         )
 
         self.Rota.build_and_solve(options={"ratio": 0.10})
 
-        self.Rota.export_rota_to_html("nwd")
+        self.Rota.export_rota_to_html("nwd1")
 
         assert self.Rota.results.solver.status == "ok"
         assert self.Rota.results.solver.termination_condition == "optimal"
@@ -86,10 +80,7 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),
@@ -131,10 +122,7 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),
@@ -192,10 +180,7 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),
@@ -234,10 +219,7 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=10,
                 workers_required=2,
                 force_as_block_unless_nwd=True,
@@ -290,20 +272,14 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=10,
                 workers_required=2,
                 force_as_block_unless_nwd=True,
                 assign_as_block=True,
             ),
             SingleShift(
-                ("group1",),
-                "w",
-                12.5,
-                days[5:],
+                sites=("group1",), name="w", length= 12.5, days=days[5:],
                 balance_offset=10,
                 workers_required=4,
                 force_as_block_unless_nwd=True,
@@ -359,20 +335,14 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=10,
                 workers_required=1,
                 force_as_block_unless_nwd=True,
                 assign_as_block=True,
             ),
             SingleShift(
-                ("group1",),
-                "w",
-                12.5,
-                days[5:],
+                sites=("group1",), name="w", length= 12.5, days=days[5:],
                 balance_offset=10,
                 workers_required=2,
                 force_as_block_unless_nwd=True,

test/test_remote.py

@@ -30,14 +30,11 @@ class TestRemoteRotas:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=1,
                 balance_offset=100,
                 force_as_block=False,
-                constraint=[("require_remote_site_presence_week", ("group2", 2))],
+                constraint=[{"name":"require_remote_site_presence_week", "options":("group2", 2)}],
             ),
         )
 
@@ -69,20 +66,14 @@ class TestRemoteRotas:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days[:3],
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[:3],
                 workers_required=1,
                 balance_offset=10,
                 force_as_block=True,
-                constraint=[("require_remote_site_presence_week", ("group2", 1))],
+                constraint=[{"name":"require_remote_site_presence_week", "options":("group2", 1)}],
             ),
             SingleShift(
-                ("group1", "group2"),
-                "b",
-                12.5,
-                days[3:],
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[3:],
                 workers_required=1,
                 balance_offset=10,
                 force_as_block=False,
@@ -135,20 +126,14 @@ class TestRemoteRotas:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days[:3],
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[:3],
                 workers_required=1,
                 balance_offset=10,
                 force_as_block=False,
-                constraint=[("require_remote_site_presence", ("group2", 1))],
+                constraint=[{"name":"require_remote_site_presence_week", "options":("group2", 1)}],
             ),
             SingleShift(
-                ("group1", "group2"),
-                "b",
-                12.5,
-                days[3:],
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[3:],
                 workers_required=1,
                 balance_offset=10,
                 force_as_block=False,

test/test_rota.py

@@ -1,4 +1,5 @@
 from copy import deepcopy
+from re import S
 from black import main
 import pytest
 from rota.shifts import NoWorkers, RotaBuilder, SingleShift, days
@@ -17,7 +18,6 @@ class TestDemoRota:
         start_date,
         weeks_to_rota=weeks_to_rota,
         balance_offset_modifier=1,
-        max_weekend_frequency=1,
     )
 
     Rota.constraint_options["max_shifts_per_week"] = 5
@@ -44,10 +44,10 @@ class TestDemoRota:
     # Add a weekday and weekend shift
     Rota.add_shifts(
         SingleShift(
-            ("group1", "group2"),
-            "weekday",
-            12.5,
-            days[:5],
+            sites=("group1", "group2"),
+            name="weekday",
+            length=12.5,
+            days=days[:5],
             balance_offset=1,
             workers_required=1,
             # We use a different balance weighting for each shift as otherwise
@@ -55,10 +55,10 @@ class TestDemoRota:
             # balance_weighting=0.6,
         ),
         SingleShift(
-            ("group1",),
-            "weekend",
-            12.5,
-            days[5:],
+            sites=("group1",),
+            name="weekend",
+            length=12.5,
+            days=days[5:],
             balance_offset=4,
             # balance_weighting=0.5
         ),
@@ -191,7 +191,6 @@ class TestDemoRotaNights:
         start_date,
         weeks_to_rota=weeks_to_rota,
         balance_offset_modifier=1,
-        max_weekend_frequency=1,
     )
 
     Rota.constraint_options["max_shifts_per_week"] = 5
@@ -223,12 +222,12 @@ class TestDemoRotaNights:
     # Add a weekday and weekend shift
     Rota.add_shifts(
         SingleShift(
-            ("group1", "group2"),
-            "night_weekday",
-            12.5,
-            days[:4],
+            sites=("group1", "group2"),
+            name="night_weekday",
+            length=12.5,
+            days=days[:4],
             workers_required=1,
-            constraint=["preclear2", "postclear2"],
+            constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
             # constraint=["night"],
             force_as_block=True
             # We use a different balance weighting for each shift as otherwise
@@ -236,20 +235,20 @@ class TestDemoRotaNights:
             # balance_weighting=0.6,
         ),
         SingleShift(
-            ("group1", "group2"),
-            "night_weekend",
-            12.5,
-            days[4:],
+            sites=("group1", "group2"),
+            name="night_weekend",
+            length=12.5,
+            days=days[4:],
             # balance_weighting=0.5
-            constraint=["preclear2", "postclear2"],
+            constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
             # constraint=["night"],
             force_as_block=True,
         ),
         SingleShift(
-            ("group1", "group2"),
-            "twilight",
-            12.5,
-            days[:5],
+            sites=("group1", "group2"),
+            name="twilight",
+            length=12.5,
+            days=days[:5],
             # balance_weighting=0.5
             workers_required=2,
             # constraint=["preclear2", "postclear2"],
@@ -271,11 +270,9 @@ class TestDemoRotaNights:
 
     Rota.export_rota_to_html("test3")
 
-
     def test_start_date(self):
         assert self.Rota.start_date == self.start_date
 
-
     def test_night_assignment(self):
         shift_summary = self.Rota.get_shift_summary_dict()
 
@@ -303,7 +300,6 @@ class TestDemoRotaClear:
         start_date,
         weeks_to_rota=weeks_to_rota,
         balance_offset_modifier=1,
-        max_weekend_frequency=1,
     )
     Rota.constraint_options["max_shifts_per_month"] = 20
 
@@ -323,30 +319,30 @@ class TestDemoRotaClear:
     def test_preclear(self):
         self.Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days[:5],
+                sites=("group1", "group2"),
+                name="a",
+                length=12.5,
+                days=days[:5],
                 balance_offset=40,
                 workers_required=2,
-                constraint=["preclear"],
+                constraint=[{"name": "preclear"}],
             ),
             SingleShift(
-                ("group1", "group2"),
-                "b",
-                12.5,
-                days[:5],
+                sites=("group1", "group2"),
+                name="b",
+                length=12.5,
+                days=days[:5],
                 balance_offset=40,
                 workers_required=2,
             ),
             SingleShift(
-                ("group1", "group2"),
-                "c",
-                12.5,
-                days[3],
+                sites=("group1", "group2"),
+                name="c",
+                length=12.5,
+                days=days[3],
                 balance_offset=40,
                 workers_required=1,
-                constraint=["preclear2", "postclear2"],
+                constraint=[{"name": "preclear2"}, {"name": "postclear2"}],
             ),
         )
 
@@ -389,7 +385,6 @@ class TestDemoRotaShiftConstraints:
         start_date,
         weeks_to_rota=weeks_to_rota,
         balance_offset_modifier=1,
-        max_weekend_frequency=1,
     )
     # Rota.constraint_options["max_shifts_per_week"] = 5
     # Rota.constraint_options["max_shifts_per_month"] = 20
@@ -399,10 +394,10 @@ class TestDemoRotaShiftConstraints:
     Rota.add_workers((worker1, worker2))
     Rota.add_shifts(
         SingleShift(
-            ("group1",),
-            "a",
-            12.5,
-            days,
+            sites=("group1",),
+            name="a",
+            length=12.5,
+            days=days,
             balance_offset=40,
             workers_required=1,
         ),
@@ -479,14 +474,14 @@ class TestDemoRotaBalanceShiftSites:
     def test_balance_blocks_across_2_groups(self):
         self.Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "weekday_night",
-                12.5,
-                days[:4],
+                sites=("group1", "group2"),
+                name="weekday_night",
+                length=12.5,
+                days=days[:4],
                 balance_offset=40,
                 workers_required=2,
                 force_as_block=True,
-                constraint=["balance_across_groups"],
+                constraint=[{"name": "balance_across_groups"}],
             ),
         )
         self.Rota.constraint_options["balance_nights_across_sites"] = False
@@ -515,14 +510,14 @@ class TestDemoRotaBalanceShiftSites:
         self.Rota.shifts = []
         self.Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "weekday_night",
-                12.5,
-                days[:4],
+                sites=("group1", "group2"),
+                name="weekday_night",
+                length=12.5,
+                days=days[:4],
                 balance_offset=40,
                 workers_required=3,
                 force_as_block=True,
-                constraint=["balance_across_groups"],
+                constraint=[{"name": "balance_across_groups"}],
             ),
         )
         self.Rota.constraint_options["balance_nights_across_sites"] = False
@@ -549,14 +544,14 @@ class TestDemoRotaBalanceShiftSites:
 
         self.Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days[:4],
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days[:4],
                 balance_offset=40,
                 workers_required=3,
                 force_as_block=True,
-                constraint=["balance_across_groups"],
+                constraint=[{"name": "balance_across_groups"}],
             ),
         )
         self.Rota.constraint_options["balance_nights_across_sites"] = False
@@ -580,8 +575,7 @@ class TestDemoRotaBalanceShiftSites:
         assert self.Rota.results.solver.status == "ok"
         assert self.Rota.results.solver.termination_condition == "optimal"
 
-
-class TestLimitConstraints:
+def generate_limit_constraint_rota():
     weeks_to_rota = 8
     start_date = datetime.date(2022, 3, 7)
 
@@ -603,30 +597,41 @@ class TestLimitConstraints:
 
     Rota.shifts = []
 
+    return Rota
+
+class TestLimitConstraints:
+
     def test_constraint_limit_grades(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days,
-                balance_offset=40,
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days,
+                balance_offset=60,
                 workers_required=4,
                 force_as_block=True,
-                constraint=[("limit_grade_number", {2: 1})],
+                constraint=[{"name": "limit_grade_number", "options": {2: 1}}],
             ),
         )
-        # self.Rota.constraint_options["balance_nights_across_sites"] = False
-        # self.Rota.constraint_options["balance_shifts_over_workers"] = False
-        self.Rota.constraint_options["balance_shifts_quadratic"] = True
-        self.Rota.build_and_solve(options={"ratio": 0.00})
+        # Rota.constraint_options["balance_nights_across_sites"] = False
+        # Rota.constraint_options["balance_shifts_over_workers"] = False
+        #Rota.constraint_options["balance_shifts_quadratic"] = True
+        Rota.build_and_solve(options={"ratio": 0.1})
+
+        Rota.export_rota_to_html("constraint_limit_grades")
+
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
+        grade_workers = Rota.get_workers_by_grade()
 
-        grade_workers = self.Rota.get_workers_by_grade()
         for grade in grade_workers:
             shift_patterns = []
             for w in grade_workers[grade]:
-                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+                shift_patterns.append(Rota.get_worker_shift_list(w))
 
             zipped_lists = list(zip(*shift_patterns))
 
@@ -636,35 +641,33 @@ class TestLimitConstraints:
             for day_shifts in zipped_lists:
                 assert day_shifts.count("weekday_night") == limit
 
-        assert self.Rota.results.solver.status == "ok"
-        assert self.Rota.results.solver.termination_condition == "optimal"
-
     def test_constraint_limit_grades2(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days,
                 balance_offset=40,
                 workers_required=4,
                 force_as_block=True,
-                constraint=[("limit_grade_number", {3: 2})],
+                constraint=[{"name": "limit_grade_number", "options": {3: 2}}],
             ),
         )
-        # self.Rota.constraint_options["balance_nights_across_sites"] = False
-        # self.Rota.constraint_options["balance_shifts_over_workers"] = False
-        self.Rota.constraint_options["balance_shifts_quadratic"] = True
-        self.Rota.build_and_solve(options={"ratio": 0.1})
+        # Rota.constraint_options["balance_nights_across_sites"] = False
+        # Rota.constraint_options["balance_shifts_over_workers"] = False
+        Rota.constraint_options["balance_shifts_quadratic"] = True
+        Rota.build_and_solve(options={"ratio": 0.1})
 
-        self.Rota.export_rota_to_html("test9")
+        Rota.export_rota_to_html("test9")
 
-        grade_workers = self.Rota.get_workers_by_grade()
+        grade_workers = Rota.get_workers_by_grade()
         for grade in grade_workers:
             shift_patterns = []
             for w in grade_workers[grade]:
-                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+                shift_patterns.append(Rota.get_worker_shift_list(w))
 
             zipped_lists = list(zip(*shift_patterns))
 
@@ -674,35 +677,51 @@ class TestLimitConstraints:
             for day_shifts in zipped_lists:
                 assert day_shifts.count("weekday_night") == limit
 
-        assert self.Rota.results.solver.status == "ok"
-        assert self.Rota.results.solver.termination_condition == "optimal"
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
 
     def test_constraint_limit_grades3(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+
+        Rota.workers = []
+
+        worker1 = Worker(name="worker1", site="group1", grade=2, remote_site="group1")
+        worker2 = Worker(name="worker2", site="group1", grade=2, remote_site="group1")
+        worker3 = Worker(name="worker3", site="group1", grade=2, remote_site="group1")
+        worker4 = Worker(name="worker4", site="group3", grade=3, remote_site="group2")
+        worker5 = Worker(name="worker5", site="group3", grade=3, remote_site="group2")
+        worker6 = Worker(name="worker6", site="group3", grade=3, remote_site="group2")
+        worker7 = Worker(name="worker7", site="group3", grade=2, remote_site="group2")
+
+        Rota.add_workers((worker1, worker2, worker3, worker4))
+        Rota.add_workers((worker5, worker6))
+        
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days,
-                balance_offset=40,
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days,
+                balance_offset=60,
                 workers_required=4,
                 force_as_block=True,
-                constraint=[("limit_grade_number", {2: 3, 3: 1})],
+                constraint=[{"name": "limit_grade_number", "options": {2: 3, 3: 1}},]
             ),
         )
-        # self.Rota.constraint_options["balance_nights_across_sites"] = False
-        # self.Rota.constraint_options["balance_shifts_over_workers"] = False
-        self.Rota.constraint_options["balance_shifts_quadratic"] = True
-        self.Rota.build_and_solve(options={"ratio": 0.1})
+        # Rota.constraint_options["balance_nights_across_sites"] = False
+        # Rota.constraint_options["balance_shifts_over_workers"] = False
+        Rota.constraint_options["balance_shifts_quadratic"] = True
+        Rota.build_and_solve(options={"ratio": 0.01})
 
-        self.Rota.export_rota_to_html("test9")
+        Rota.export_rota_to_html("test9")
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
 
-        grade_workers = self.Rota.get_workers_by_grade()
+        grade_workers = Rota.get_workers_by_grade()
         for grade in grade_workers:
             shift_patterns = []
             for w in grade_workers[grade]:
-                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+                shift_patterns.append(Rota.get_worker_shift_list(w))
 
             zipped_lists = list(zip(*shift_patterns))
 
@@ -712,55 +731,57 @@ class TestLimitConstraints:
             for day_shifts in zipped_lists:
                 assert day_shifts.count("weekday_night") == limit
 
-        assert self.Rota.results.solver.status == "ok"
-        assert self.Rota.results.solver.termination_condition == "optimal"
 
     def test_constraint_limit_grades4(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days,
                 balance_offset=40,
                 workers_required=4,
                 force_as_block=True,
-                constraint=[("limit_grade_number", {2: 4, 3: 0})],
+                constraint=[{"name": "limit_grade_number", "options": {2: 4, 3: 0}}],
             ),
         )
-        # self.Rota.constraint_options["balance_nights_across_sites"] = False
-        # self.Rota.constraint_options["balance_shifts_over_workers"] = False
-        self.Rota.constraint_options["balance_shifts_quadratic"] = True
-        self.Rota.build_and_solve(options={"ratio": 0.1})
+        # Rota.constraint_options["balance_nights_across_sites"] = False
+        # Rota.constraint_options["balance_shifts_over_workers"] = False
+        Rota.constraint_options["balance_shifts_quadratic"] = True
+        Rota.build_and_solve(options={"ratio": 0.1})
 
-        assert self.Rota.results.solver.termination_condition == "infeasible"
+        assert Rota.results.solver.termination_condition == "infeasible"
 
     def test_constraint_minimum_grades(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days,
                 balance_offset=40,
                 workers_required=1,
                 force_as_block=True,
-                constraint=[("minimum_grade_number", (3, 1))],
+                constraint=[{"name": "minimum_grade_number", "options": (3, 1)}],
             ),
         )
 
-        self.Rota.constraint_options["balance_shifts_quadratic"] = True
-        self.Rota.build_and_solve(options={"ratio": 0.1})
+        Rota.constraint_options["balance_shifts_quadratic"] = True
+        Rota.build_and_solve(options={"ratio": 0.1})
 
-        self.Rota.export_rota_to_html("test9")
+        Rota.export_rota_to_html("test9")
 
-        grade_workers = self.Rota.get_workers_by_grade()
+        assert Rota.results.solver.status == "ok"
+
+        grade_workers = Rota.get_workers_by_grade()
         for grade in grade_workers:
             shift_patterns = []
             for w in grade_workers[grade]:
-                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+                shift_patterns.append(Rota.get_worker_shift_list(w))
 
             zipped_lists = list(zip(*shift_patterns))
 
@@ -770,33 +791,33 @@ class TestLimitConstraints:
             for day_shifts in zipped_lists:
                 assert day_shifts.count("weekday_night") == limit
 
-        assert self.Rota.results.solver.status == "ok"
 
     def test_constraint_minimum_grades2(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days,
                 balance_offset=40,
                 workers_required=4,
                 force_as_block=True,
-                constraint=[("minimum_grade_number", (3, 3))],
+                constraint=[{"name": "minimum_grade_number", "options": (3, 3)}],
             ),
         )
 
-        self.Rota.constraint_options["balance_shifts_quadratic"] = True
-        self.Rota.build_and_solve(options={"ratio": 0.1})
+        Rota.constraint_options["balance_shifts_quadratic"] = True
+        Rota.build_and_solve(options={"ratio": 0.1})
 
-        self.Rota.export_rota_to_html("test9")
+        Rota.export_rota_to_html("test9")
 
-        grade_workers = self.Rota.get_workers_by_grade()
+        grade_workers = Rota.get_workers_by_grade()
         for grade in grade_workers:
             shift_patterns = []
             for w in grade_workers[grade]:
-                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+                shift_patterns.append(Rota.get_worker_shift_list(w))
 
             zipped_lists = list(zip(*shift_patterns))
 
@@ -806,71 +827,79 @@ class TestLimitConstraints:
             for day_shifts in zipped_lists:
                 assert day_shifts.count("weekday_night") == limit
 
-        assert self.Rota.results.solver.status == "ok"
+        assert Rota.results.solver.status == "ok"
 
     def test_constraint_minimum_grades3(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days,
                 balance_offset=40,
                 workers_required=4,
                 force_as_block=True,
-                constraint=[("minimum_grade_number", (3, 0))],
+                constraint=[{"name": "minimum_grade_number", "options": (3, 0)}],
             ),
         )
 
-        self.Rota.constraint_options["balance_shifts_quadratic"] = True
-        self.Rota.build_and_solve(options={"ratio": 0.1})
+        Rota.constraint_options["balance_shifts_quadratic"] = True
+        Rota.build_and_solve(options={"ratio": 0.1})
 
-        self.Rota.export_rota_to_html("test9")
+        Rota.export_rota_to_html("test9")
 
-        assert self.Rota.results.solver.status == "ok"
+        assert Rota.results.solver.status == "ok"
 
     def test_constraint_minimum_grades_no_valid_worker(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "weekday_night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="weekday_night",
+                length=12.5,
+                days=days,
                 balance_offset=40,
                 workers_required=4,
                 force_as_block=True,
-                constraint=[("minimum_grade_number", (4, 1))],
+                constraint=[{"name": "minimum_grade_number", "options": (4, 1)}],
             ),
         )
 
-        self.Rota.constraint_options["balance_shifts_quadratic"] = True
+        Rota.constraint_options["balance_shifts_quadratic"] = True
         with pytest.raises(ValueError):
-            self.Rota.build_and_solve(options={"ratio": 0.1})
+            Rota.build_and_solve(options={"ratio": 0.1})
 
     def test_constraint_require_remote_site_presence_week(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="night",
+                length=12.5,
+                days=days,
                 balance_offset=40,
                 workers_required=2,
                 force_as_block=True,
-                constraint=[("require_remote_site_presence_week", ("group1", 2))],
+                constraint=[
+                    {
+                        "name": "require_remote_site_presence_week",
+                        "options": ("group1", 2),
+                    }
+                ],
             ),
         )
-        self.Rota.build_and_solve(options={"ratio": 0.00})
-        self.Rota.export_rota_to_html("remote1")
+        Rota.build_and_solve(options={"ratio": 0.00})
+        Rota.export_rota_to_html("remote1")
 
-        group_workers = self.Rota.get_workers_by_group()
+        group_workers = Rota.get_workers_by_group()
         for group in group_workers:
             shift_patterns = []
             for w in group_workers[group]:
-                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+                shift_patterns.append(Rota.get_worker_shift_list(w))
 
             zipped_lists = list(zip(*shift_patterns))
 
@@ -880,67 +909,79 @@ class TestLimitConstraints:
             for day_shifts in zipped_lists:
                 assert day_shifts.count("night") == limit
 
-        assert self.Rota.results.solver.status == "ok"
-        assert self.Rota.results.solver.termination_condition == "optimal"
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
 
     def test_constraint_require_remote_site_presence_week2(self):
-        self.Rota.constraint_options["minimise_shift_diffs"] = False
-        self.Rota.constraint_options["balance_shifts_quadratic"] = False
-        self.Rota.constraint_options["balance_shift"] = False
-        self.Rota.constraint_options["balance_nights_across_sites"] = True
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.constraint_options["minimise_shift_diffs"] = False
+        Rota.constraint_options["balance_shifts_quadratic"] = False
+        Rota.constraint_options["balance_shift"] = False
+        Rota.constraint_options["balance_nights_across_sites"] = True
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="night",
+                length=12.5,
+                days=days,
                 balance_offset=20,
                 workers_required=3,
                 force_as_block=True,
-                constraint=[("require_remote_site_presence_week", ("group2", 2))],
+                constraint=[
+                    {
+                        "name": "require_remote_site_presence_week",
+                        "options": ("group1", 2),
+                    }
+                ],
             ),
         )
-        self.Rota.build_and_solve(options={"ratio": 0.00})
-        self.Rota.export_rota_to_html("remote2")
+        Rota.build_and_solve(options={"ratio": 0.00})
+        Rota.export_rota_to_html("remote2")
 
-        group_workers = self.Rota.get_workers_by_remote_group()
+        group_workers = Rota.get_workers_by_remote_group()
         for group in group_workers:
             shift_patterns = []
             for w in group_workers[group]:
-                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+                shift_patterns.append(Rota.get_worker_shift_list(w))
 
             zipped_lists = list(zip(*shift_patterns))
 
-            if group == "group2":
+            if group == "group1":
                 for day_shifts in zipped_lists:
                     assert day_shifts.count("night") >= 2
 
-        assert self.Rota.results.solver.status == "ok"
-        assert self.Rota.results.solver.termination_condition == "optimal"
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
 
     def test_constraint_require_remote_site_presence_week3(self):
-        self.Rota.shifts = []
-        self.Rota.add_shifts(
+        Rota = generate_limit_constraint_rota()
+        Rota.shifts = []
+        Rota.add_shifts(
             SingleShift(
-                ("group1", "group2", "group3"),
-                "night",
-                12.5,
-                days,
+                sites=("group1", "group2", "group3"),
+                name="night",
+                length=12.5,
+                days=days,
                 balance_offset=20,
                 workers_required=3,
                 force_as_block=True,
-                constraint=[("require_remote_site_presence_week", ("group2", 1))],
+                constraint=[
+                    {
+                        "name": "require_remote_site_presence_week",
+                        "options": ("group2", 1),
+                    }
+                ],
             ),
         )
-        self.Rota.build_and_solve(options={"ratio": 0.00})
-        self.Rota.export_rota_to_html("remote3")
+        Rota.build_and_solve(options={"ratio": 0.00})
+        Rota.export_rota_to_html("remote3")
 
-        group_workers = self.Rota.get_workers_by_remote_group()
+        group_workers = Rota.get_workers_by_remote_group()
         for group in group_workers:
             shift_patterns = []
             for w in group_workers[group]:
-                shift_patterns.append(self.Rota.get_worker_shift_list(w))
+                shift_patterns.append(Rota.get_worker_shift_list(w))
 
             zipped_lists = list(zip(*shift_patterns))
 
@@ -948,8 +989,9 @@ class TestLimitConstraints:
                 for day_shifts in zipped_lists:
                     assert day_shifts.count("night") >= 1
 
-        assert self.Rota.results.solver.status == "ok"
-        assert self.Rota.results.solver.termination_condition == "optimal"
+        assert Rota.results.solver.status == "ok"
+        assert Rota.results.solver.termination_condition == "optimal"
+
 
 # class TestNoWorkerRota:
 #    weeks_to_rota = 10
@@ -979,23 +1021,41 @@ class TestNightUnavailable:
 
     # Add a few workers
     worker1 = Worker(
-        name="worker1", site="group1", grade=1, not_available_to_work=[{ "date": datetime.datetime.strptime("15/03/22", "%d/%m/%y").date(), "reason" : "****" },]
+        name="worker1",
+        site="group1",
+        grade=1,
+        not_available_to_work=[
+            {
+                "date": datetime.datetime.strptime("15/03/22", "%d/%m/%y").date(),
+                "reason": "****",
+            },
+        ],
     )
     worker2 = Worker(
-        name="worker2", site="group1", grade=1, not_available_to_work=[ { "date": datetime.datetime.strptime("14/03/22", "%d/%m/%y").date(), "reason": "****" } ]
+        name="worker2",
+        site="group1",
+        grade=1,
+        not_available_to_work=[
+            {
+                "date": datetime.datetime.strptime("14/03/22", "%d/%m/%y").date(),
+                "reason": "****",
+            }
+        ],
     )
 
     Rota.add_workers((worker1, worker2))
 
-
-
     def test_basic_assignment(self, capsys):
         self.Rota.shifts = []
 
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",), "night_weekend", 12.5, days[5:], constraint=["night"]
+                sites =("group1",),
+                name = "night_weekend",
+                length = 12.5,
+                days = days[5:],
+                constraint=[{"name": "night"}],
             ),
         )
 
@@ -1005,7 +1065,6 @@ class TestNightUnavailable:
         assert self.Rota.results.solver.status == "ok"
         assert self.Rota.results.solver.termination_condition == "optimal"
 
-
         self.Rota.export_rota_to_html("testnight")
         summary = self.Rota.get_shift_summary_dict()
 
@@ -1018,7 +1077,12 @@ class TestNightUnavailable:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",), "night_weekend", 12.5, days[5:], workers_required=2, constraint=["night"]
+                sites =("group1",),
+                name = "night_weekend",
+                length = 12.5,
+                days = days[5:],
+                workers_required=2,
+                constraint=[{"name": "night"}],
             ),
         )
 
@@ -1032,7 +1096,12 @@ class TestNightUnavailable:
 
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",), "night_weekend", 12.5, days[5:], workers_required=2, constraint=[]
+                sites = ("group1",),
+                name = "night_weekend",
+                length = 12.5,
+                days = days[5:],
+                workers_required=2,
+                constraint=[],
             ),
         )
 
@@ -1044,16 +1113,31 @@ class TestNightUnavailable:
     def test_assign_split(self):
         self.Rota.shifts = []
         worker3 = Worker(
-            name="worker3", site="group1", grade=1, not_available_to_work=({"date":datetime.datetime.strptime("13/03/22", "%d/%m/%y").date()},)
+            name="worker3",
+            site="group1",
+            grade=1,
+            not_available_to_work=(
+                {"date": datetime.datetime.strptime("13/03/22", "%d/%m/%y").date()},
+            ),
         )
         worker4 = Worker(
-            name="worker4", site="group1", grade=1, not_available_to_work=({"date":datetime.datetime.strptime("12/03/22", "%d/%m/%y").date()},)
+            name="worker4",
+            site="group1",
+            grade=1,
+            not_available_to_work=(
+                {"date": datetime.datetime.strptime("12/03/22", "%d/%m/%y").date()},
+            ),
         )
         self.Rota.add_workers((worker3, worker4))
 
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",), "night_weekend", 12.5, days[5:], workers_required=3, constraint=[]
+                sites=("group1",),
+                name="night_weekend",
+                length=12.5,
+                days=days[5:],
+                workers_required=3,
+                constraint=[],
             ),
         )
 
@@ -1067,6 +1151,7 @@ class TestNightUnavailable:
         for worker in summary:
             assert summary[worker]["night_weekend"] == 15
 
+
 #    def test_assign_split_night_constraint(self):
 #        self.Rota.shifts = []
 #        worker3 = Worker(

test/test_shifts.py

@@ -49,10 +49,7 @@ class TestShifts:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
             ),
         )
 
@@ -63,12 +60,116 @@ class TestShifts:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1"),
-                "a",
-                12.5,
-                days,
+                sites=("group1",), name="a", length= 12.5, days=days,
             ),
         )
 
         with pytest.raises(InvalidShift):
             Rota.build_and_solve(options={"ratio": 0.000})
+
+
+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,
+                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):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            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.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(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            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(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            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 == "warning"
+
+    def test_max_shifts_per_week_block_shift_number(self):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            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):
+        Rota = generate_basic_rota()
+
+        Rota.add_shifts(
+            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 == "warning"

test/test_weekend_padding.py

@@ -33,18 +33,12 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "a",
-                12.5,
-                days,
+                sites=("group1",), name="a", length= 12.5, days=days,
                 balance_offset=1,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "b",
-                12.5,
-                days[5:],
+                sites=("group1",), name="b", length= 12.5, days=days[5:],
                 balance_offset=0,
                 workers_required=1
             ),
@@ -97,19 +91,13 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "a",
-                12.5,
-                days,
+                sites=("group1",), name="a", length= 12.5, days=days,
                 balance_offset=5,
                 workers_required=1,
                 force_as_block=True
             ),
             SingleShift(
-                ("group1",),
-                "b",
-                12.5,
-                days[5:],
+                sites=("group1",), name="b", length= 12.5, days=days[5:],
                 balance_offset=0,
                 workers_required=2
             ),
@@ -159,18 +147,12 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "w",
-                12.5,
-                days[5:],
+                sites=("group1",), name="w", length= 12.5, days=days[5:],
                 workers_required=1
             ),
         )
@@ -211,26 +193,17 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "c",
-                12.5,
-                days[1],
+                sites=("group1",), name="c", length= 12.5, days=days[1],
                 balance_offset=1,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "w",
-                12.5,
-                days[5:],
+                sites=("group1",), name="w", length= 12.5, days=days[5:],
                 workers_required=1
             ),
         )
@@ -259,26 +232,17 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=1,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "c",
-                12.5,
-                days[1],
+                sites=("group1",), name="c", length= 12.5, days=days[1],
                 balance_offset=1,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "w",
-                12.5,
-                days[5:],
+                sites=("group1",), name="w", length= 12.5, days=days[5:],
                 workers_required=1
             ),
         )
@@ -318,26 +282,17 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=1,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "c",
-                12.5,
-                days[0],
+                sites=("group1",), name="c", length= 12.5, days=days[0],
                 balance_offset=1,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "w",
-                12.5,
-                days[5:],
+                sites=("group1",), name="w", length= 12.5, days=days[5:],
                 workers_required=1
             ),
         )
@@ -368,26 +323,17 @@ class TestDemoRota:
         # Add a weekday and weekend shift
         self.Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",), name="d", length= 12.5, days=days[:5],
                 balance_offset=1,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "c",
-                12.5,
-                days[0],
+                sites=("group1",), name="c", length= 12.5, days=days[0],
                 balance_offset=1,
                 workers_required=1,
             ),
             SingleShift(
-                ("group1",),
-                "w",
-                12.5,
-                days[5:],
+                sites=("group1",), name="w", length= 12.5, days=days[5:],
                 workers_required=1,
                 force_as_block=True
             ),

test/test_worker_requests.py

@@ -50,10 +50,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=2,
                 force_as_block=False,
             ),
@@ -90,10 +87,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=2,
                 force_as_block=False,
             ),
@@ -128,10 +122,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=2,
                 force_as_block=False,
             ),
@@ -173,10 +164,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=2,
                 force_as_block=False,
             ),
@@ -218,10 +206,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=2,
                 balance_offset=1,
                 force_as_block=False,
@@ -255,10 +240,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=2,
                 balance_offset=1,
                 force_as_block=False,
@@ -312,10 +294,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=3,
                 balance_offset=1,
                 force_as_block=False,
@@ -356,10 +335,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=2,
                 balance_offset=1,
                 force_as_block=False,
@@ -401,10 +377,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "a",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
                 workers_required=1,
                 balance_offset=1,
                 force_as_block=False,
@@ -417,10 +390,7 @@ class TestWorkerRequests:
         
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "c",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="c", length= 12.5, days=days,
                 workers_required=1,
                 balance_offset=1,
                 force_as_block=False,
@@ -432,10 +402,7 @@ class TestWorkerRequests:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "b",
-                12.5,
-                days,
+                sites=("group1", "group2"), name="b", length= 12.5, days=days,
                 workers_required=1,
                 balance_offset=1,
                 force_as_block=False,
@@ -461,3 +428,104 @@ class TestWorkerRequests:
 
         with pytest.raises(InvalidShift):
             Rota.build_and_solve(options={"ratio": 0.000})
+
+    def test_work_requests_generic(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+        Rota.constraint_options["balance_bank_holidays"] = False
+
+        Rota.add_workers(
+            [Worker(
+                name="worker3", site="group1", grade=1,
+                work_requests=(
+                    [
+                        {"date": i, "shift": "*"}
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
+                    ]
+                ),
+            ),
+            Worker(
+                name="worker4", site="group1", grade=1,
+                work_requests=(
+                    [
+                        {"date": i, "shift": "*"}
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 35))
+                    ]
+                ),
+            ),]
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days,
+                workers_required=2,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.000})
+        Rota.export_rota_to_html("work_requests")
+
+        shift_summary = Rota.get_shift_summary_dict()
+        for worker_name in shift_summary:
+            worker_shifts = shift_summary[worker_name]
+
+            worker = Rota.get_worker_by_name(worker_name)
+
+            assert worker_shifts["a"] == 35
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            if worker.name in ("worker3", "worker4"):
+               assert shift_string.startswith("a"*35)
+
+    def test_work_requests_generic2(self):
+        """_summary_"""
+        Rota = generate_basic_rota()
+        Rota.constraint_options["balance_bank_holidays"] = False
+
+        Rota.add_workers(
+            [Worker(
+                name="worker3", site="group1", grade=1,
+                work_requests=(
+                    [
+                        {"date": i, "shift": "*"}
+                        for i in list(date_generator(datetime.date(2022, 3, 7), 20))
+                    ]
+                ),
+            ),
+            ]
+        )
+
+        Rota.add_shifts(
+            SingleShift(
+                sites=("group1", "group2"), name="a", length= 12.5, days=days[3:],
+                workers_required=2,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+            SingleShift(
+                sites=("group1", "group2"), name="b", length= 12.5, days=days[:3],
+                workers_required=1,
+                balance_offset=1,
+                force_as_block=False,
+            ),
+        )
+
+        Rota.build_and_solve(options={"ratio": 0.1})
+        Rota.export_rota_to_html("work_requests")
+
+        shift_summary = Rota.get_shift_summary_dict()
+        for worker_name in shift_summary:
+            worker_shifts = shift_summary[worker_name]
+
+            worker = Rota.get_worker_by_name(worker_name)
+
+            assert worker_shifts["a"] in (26, 27)
+            assert worker_shifts["b"] == 10
+
+            shift_string = Rota.get_worker_shift_list_string(worker)
+
+            if worker.name in ("worker3",):
+               assert "-" not in shift_string[:20]

test/test_workers.py

@@ -22,14 +22,20 @@ class TestWorkers:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",),
+                name="d",
+                length=12.5,
+                days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),
-            SingleShift(("group1",), "w", 12.5, days[5:], workers_required=1),
+            SingleShift(
+                sites=("group1",),
+                name="w",
+                length=12.5,
+                days=days[5:],
+                workers_required=1,
+            ),
         )
 
         with pytest.raises(NoWorkers):
@@ -49,10 +55,10 @@ class TestWorkers:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1",),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1",),
+                name="d",
+                length=12.5,
+                days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),
@@ -128,10 +134,10 @@ class TestWorkers:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1", "group2"),
+                name="d",
+                length=12.5,
+                days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),
@@ -206,7 +212,8 @@ class TestWorkers:
         worker3_oop = [
             {
                 "start_date": start_date + datetime.timedelta(weeks=weeks_to_rota / 2),
-                "end_date": start_date + datetime.timedelta(weeks=3 * weeks_to_rota / 4),
+                "end_date": start_date
+                + datetime.timedelta(weeks=3 * weeks_to_rota / 4),
                 "reason": "Quarter oop",
             }
         ]
@@ -282,10 +289,10 @@ class TestWorkers:
 
         Rota.add_shifts(
             SingleShift(
-                ("group1", "group2"),
-                "d",
-                12.5,
-                days[:5],
+                sites=("group1", "group2"),
+                name="d",
+                length=12.5,
+                days=days[:5],
                 balance_offset=10,
                 workers_required=1,
             ),

web/main.py

@@ -0,0 +1,63 @@
+import uvicorn
+from fastapi import FastAPI, Query
+from pydantic import BaseModel
+
+from fastapi_sqlalchemy import DBSessionMiddleware, db
+
+import os
+from dotenv import load_dotenv
+
+from rota.shifts import SingleShift, ShiftConstraint
+
+from web.models import SingleShift as ModelSingleShift
+from web.models import ShiftConstraint as ModelShiftConstraint
+
+#from rota.wokers
+
+
+load_dotenv('../.env')
+
+app = FastAPI()
+
+app.add_middleware(DBSessionMiddleware, db_url="sqlite:///test.db")
+
+@app.post('/shifts/', response_model=SingleShift)
+async def shifts(shift: ModelSingleShift):
+    db_shift = ModelSingleShift(name=shift.name, length=shift.length)
+    db.session.add(db_shift)
+    db.session.commit()
+    return db_shift
+
+    
+@app.get("/shifts/")
+async def shifts():
+    shifts = db.session.query(ModelSingleShift).all()
+
+    return shifts
+
+
+@app.get("/")
+async def root():
+    return {"message": "Hello World"}
+
+@app.get("/files/{file_path:path}")
+async def read_file(file_path: str):
+    return {"file_path": file_path}
+
+
+
+class Item(BaseModel):
+    name: str
+    description: str | None = None
+    price: float
+    tax: float | None = None
+
+
+
+@app.post("/items/")
+async def create_item(item: Item):
+    item.name + item.price
+    return item
+
+if __name__ == '__main__':
+    uvicorn.run(app, host='0.0.0.0', port=8008)

web/models.py

@@ -0,0 +1,31 @@
+from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Float, Boolean
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import relationship
+from sqlalchemy.sql import func
+
+Base = declarative_base()
+
+class SingleShift(Base):
+    __tablename__ = 'single_shift'
+    id  = Column(Integer, primary_key=True, index=True)
+    name = Column(String)
+    length = Column(Float)
+    balance_offset = Column(Float, nullable=True)
+    balance_weighting = Column(Float)
+    workers_required = Column(Float)
+    assign_as_block = Column(Boolean)
+    force_as_block = Column(Boolean)
+    hard_constrain_shift = Column(Boolean)
+    bank_holidays_only = Column(Boolean)
+    constraint_id = Column(Integer, ForeignKey('shift_constraint.id'))
+
+    constraint = relationship('ShiftConstraint')
+
+
+class ShiftConstraint(Base):
+    __tablename__ = 'shift_constraint'
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+    options = Column(Integer, nullable=True)
+
+

web/test.py

@@ -0,0 +1,54 @@
+from datetime import datetime, time, timedelta
+from uuid import UUID
+
+from fastapi import Body, FastAPI,Form, File, UploadFile
+from fastapi.responses import HTMLResponse
+
+app = FastAPI()
+
+
+@app.put("/items/{item_id}")
+async def read_items(
+    item_id: UUID,
+    start_datetime: datetime | None = Body(default=None),
+    end_datetime: datetime | None = Body(default=None),
+    repeat_at: time | None = Body(default=None),
+    process_after: timedelta | None = Body(default=None),
+):
+    start_process = start_datetime + process_after
+    duration = end_datetime - start_process
+    return {
+        "item_id": item_id,
+        "start_datetime": start_datetime,
+        "end_datetime": end_datetime,
+        "repeat_at": repeat_at,
+        "process_after": process_after,
+        "start_process": start_process,
+        "duration": duration,
+    }
+
+@app.post("/files/")
+async def create_files(files: list[bytes] = File()):
+    return {"file_sizes": [len(file) for file in files]}
+
+
+@app.post("/uploadfiles/")
+async def create_upload_files(files: list[UploadFile]):
+    return {"filenames": [file.filename for file in files]}
+
+
+@app.get("/")
+async def main():
+    content = """
+<body>
+<form action="/files/" enctype="multipart/form-data" method="post">
+<input name="files" type="file" multiple>
+<input type="submit">
+</form>
+<form action="/uploadfiles/" enctype="multipart/form-data" method="post">
+<input name="files" type="file" multiple>
+<input type="submit">
+</form>
+</body>
+    """
+    return HTMLResponse(content=content)