pre repalce

rota/shifts.py

@@ -3,6 +3,8 @@ import datetime
 import itertools
 from typing import Dict, Iterable, List, Sequence, Tuple, Set
 
+from pydantic import BaseModel, Extra
+
 from pathlib import Path
 
 import datetime
@@ -45,7 +47,7 @@ SHIFT_BOUNDS = {
 }
 
 
-class SingleShift(object):
+class SingleShift(BaseModel):
     """Class to hold all details for a shift
 
 
@@ -68,56 +70,49 @@ class SingleShift(object):
 
 
     """
+    sites: Iterable[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] = [],
 
-    def __init__(
-        self,
-        sites: Iterable[str],
-        name: str,
-        length: float,
-        shift_days: Sequence[str],
-        # balance_by_site=True,
-        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,
-        constraints: Iterable[str] = [],
-    ):
-        self.site = sites
-        self.name = name
-        self.length = length
-        self.shift_days = shift_days
+    class Config:
+        extra = Extra.allow
+
+    def __init__(self, **data: Any):
+        super().__init__(**data)
+        #self.site = sites
+        #self.name = name
+        #self.length = length
+        #self.shift_days = shift_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 balance_offset is None:
-            self.balance_offset = len(shift_days)
-        else:
-            self.balance_offset = balance_offset
 
-        # weight the shift for balancing, default is 1
-        self.balance_weighting = balance_weighting
-        self.hard_constrain_shift = hard_constrain_shift
 
-        self.workers_required = workers_required
-        self.rota_on_nwds = rota_on_nwds
-        self.force_as_block = force_as_block
-        self.force_as_block_unless_nwd = force_as_block_unless_nwd
+        if self.balance_offset is None:
+            self.balance_offset = len(self.shift_days)
 
-        if force_as_block_unless_nwd:
+
+        if self.force_as_block_unless_nwd:
             self.assign_as_block = True
 
-        self.assign_as_block = assign_as_block
+        #self.assign_as_block = assign_as_block
 
         self.constraint_options = {}
         self.constraints = []
-        for c in constraints:
+        for c in self.constraint:
             match c:
                 case (constraint, options):
                     self.constraints.append(constraint)
@@ -125,9 +120,6 @@ class SingleShift(object):
                 case constraint:
                     self.constraints.append(constraint)
 
-        self.bank_holidays_only = bank_holidays_only
-
-        # self.total_shifts = length * len(days) * len(shift_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}"
@@ -154,6 +146,7 @@ class RotaBuilder(object):
         use_previous_shifts: bool = False,
         use_shift_balance_extra: bool = False,
         use_bank_holiday_extra: bool = False,
+        SHIFT_BOUNDS = SHIFT_BOUNDS
     ):
 
         print(f"Start time: {datetime.datetime.now()}")
@@ -161,6 +154,7 @@ class RotaBuilder(object):
         print(f"Use previous shifts {use_previous_shifts}")
         print(f"Use previous shifts (balance_extra) {use_shift_balance_extra}")
 
+        self.SHIFT_BOUNDS = SHIFT_BOUNDS
         self.shifts = []  # type List[SingleShift]
 
         self.days = days
@@ -341,7 +335,7 @@ class RotaBuilder(object):
                 ((worker.id) for worker in self.workers),
                 within=NonNegativeIntegers,
                 initialize=0,
-                bounds=SHIFT_BOUNDS["bank_holiday"],
+                bounds=self.SHIFT_BOUNDS["bank_holiday"],
             )
             self.model.bank_holiday_count_w = Var(
                 ((worker.id) for worker in self.workers),
@@ -428,7 +422,7 @@ class RotaBuilder(object):
             ),
             domain=NonNegativeReals,
             initialize=0,
-            bounds=SHIFT_BOUNDS["shift_count"],
+            bounds=self.SHIFT_BOUNDS["shift_count"],
         )
 
         self.model.shift_count_diff = Var(
@@ -501,7 +495,7 @@ class RotaBuilder(object):
                 ((worker.id) for worker in self.workers),
                 domain=NonNegativeReals,
                 initialize=0,
-                bounds=SHIFT_BOUNDS["night_shift_count"],
+                bounds=self.SHIFT_BOUNDS["night_shift_count"],
             )
 
             self.model.night_shift_count_t1 = Var(
@@ -565,7 +559,7 @@ class RotaBuilder(object):
             ((worker.id) for worker in self.workers),
             domain=NonNegativeIntegers,
             initialize=0,
-            bounds=SHIFT_BOUNDS["weekend_count"],
+            bounds=self.SHIFT_BOUNDS["weekend_count"],
         )
 
         if self.constraint_options["balance_weekends"]: