From 069da08f5ad6314b475ba6424f26b73d61156acb Mon Sep 17 00:00:00 2001 From: Ross Date: Tue, 9 Dec 2025 14:09:41 +0000 Subject: [PATCH] start switching constraint options --- rota_generator/shifts.py | 281 +++++++++++++++++---------------------- 1 file changed, 124 insertions(+), 157 deletions(-) diff --git a/rota_generator/shifts.py b/rota_generator/shifts.py index 74be937..382d377 100644 --- a/rota_generator/shifts.py +++ b/rota_generator/shifts.py @@ -321,49 +321,51 @@ class SingleShift(BaseModel): return self.name[0] -class RotaConstraintOptions: - """Class to hold rota-level constraint options""" +class RotaConstraintOptions(BaseModel): + """Pydantic model holding rota-level constraint options with types, defaults and descriptions. - def __init__(self): - # self.options: { name : (default_value, type, description) } - # For each option we keep a tuple: (default_value, python_type, description) - # The types are informational and can be used when generating forms/docs. - self.constraint_options = { - "balance_nights": (True, bool, "Balance night shifts across workers"), - "constrain_time_off_after_nights": (False, bool, "Apply time-off constraints after night shifts"), - "balance_nights_across_sites": (True, bool, "Balance night shifts across sites"), - "balance_bank_holidays": (True, bool, "Balance bank-holiday shifts across staff"), - "balance_blocks": (True, bool, "Balance shifts across week-blocks"), - "balance_shifts": (True, bool, "Balance shifts across workers"), - "balance_shifts_true_quadratic": (False, bool, "Use true quadratic balancing for shifts"), - "balance_shifts_quadratic": (False, bool, "Use quadratic penalty to discourage spread across shifts"), - "balance_shifts_over_workers": (True, bool, "Balance shift distribution over workers"), - "minimise_shift_diffs": (False, bool, "Use simpler minimisation to reduce shift differences"), - "maximum_allowed_shift_diff": (None, Optional[int], "Maximum allowed per-worker shift difference (None = no limit)"), - "balance_weekends": (True, bool, "Balance weekend allocations"), - "max_weekends": (100, int, "Maximum number of weekends considered"), - "max_weekend_frequency": (1, int, "Don't assign multiple shifts every N weeks (requires balance_weekends)"), - "max_night_frequency": (1, int, "Don't assign multiple night shifts every N weeks"), - "max_night_frequency_week_exclusions": ([], list, "Weeks to exclude from night frequency checks"), - "max_shifts_per_week": (7, int, "Max shifts per worker per week"), - "max_shifts_per_month": (40, int, "Max shifts per worker per month"), - "max_days_per_week_block": ([], list, "Optional list of (max_days, week_block) tuples to limit days in blocks"), - "prevent_monday_after_full_weekends": ([], list, "Prevent assignment on Monday after full weekends"), - "prevent_monday_and_tuesday_after_full_weekends": ([], list, "Prevent Monday/Tuesday after full weekends"), - "prevent_fridays_before_full_weekends": ([], list, "Prevent Fridays before full weekends"), - "prevent_thursdays_before_full_weekends": ([], list, "Prevent Thursdays before full weekends"), - "hard_constrain_pair_separation": (False, bool, "Enforce hard pair separation constraint"), - "avoid_shifts_by_grades": ([], list, "List of (grade, shift_name) pairs or rules to avoid assigning a grade to a shift"), - "avoid_shifts_by_worker_names": ([], list, "List of worker names (or patterns) to avoid assigning to particular shifts"), - "allocate_locum_shifts": (True, bool, "Allow allocation of locum shifts"), - "distribute_locum_shifts": (True, bool, "Distribute locum shifts among available workers"), - "balance_locum_shifts": (False, bool, "Balance locum shifts similarly to normal shifts"), - "maximum_allowed_locum_shifts_per_worker": (None, Optional[int], "Maximum locum shifts allowed per worker (None = no limit)"), - "minimum_allowed_locum_shifts_per_worker": (None, Optional[int], "Minimum locum shifts required per worker (None = no minimum)"), - # This is a typed list of Pydantic constraint models; items should follow - # the schema of MinSummedGradeByShiftsPerDayConstraint - "min_summed_grade_by_shifts_per_day": ([], List[MinSummedGradeByShiftsPerDayConstraint], "List of minimum-summed-grade constraints per day"), - } + Validate/coerce raw JSON (e.g. from DB) via `RotaConstraintOptions.model_validate(raw)`. + Use `model_dump()` to produce a JSON-serializable dict for storage. + """ + + model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True) + + balance_nights: bool = Field(True, description="Balance night shifts across workers") + constrain_time_off_after_nights: bool = Field(False, description="Apply time-off constraints after night shifts") + balance_nights_across_sites: bool = Field(True, description="Balance night shifts across sites") + balance_bank_holidays: bool = Field(True, description="Balance bank-holiday shifts across staff") + balance_blocks: bool = Field(True, description="Balance shifts across week-blocks") + balance_shifts: bool = Field(True, description="Balance shifts across workers") + balance_shifts_true_quadratic: bool = Field(False, description="Use true quadratic balancing for shifts") + balance_shifts_quadratic: bool = Field(False, description="Use quadratic penalty to discourage spread across shifts") + balance_shifts_over_workers: bool = Field(True, description="Balance shift distribution over workers") + minimise_shift_diffs: bool = Field(False, description="Use simpler minimisation to reduce shift differences") + maximum_allowed_shift_diff: Optional[int] = Field(None, description="Maximum allowed per-worker shift difference (None = no limit)") + balance_weekends: bool = Field(True, description="Balance weekend allocations") + max_weekends: int = Field(100, description="Maximum number of weekends considered") + max_weekend_frequency: int = Field(1, description="Don't assign multiple shifts every N weeks (requires balance_weekends)") + max_night_frequency: int = Field(1, description="Don't assign multiple night shifts every N weeks") + max_night_frequency_week_exclusions: List[int] = Field(default_factory=list, description="Weeks to exclude from night frequency checks") + max_shifts_per_week: int = Field(7, description="Max shifts per worker per week") + max_shifts_per_month: int = Field(40, description="Max shifts per worker per month") + max_days_per_week_block: List[tuple] = Field(default_factory=list, description="Optional list of (max_days, week_block) tuples to limit days in blocks") + prevent_monday_after_full_weekends: List[int] = Field(default_factory=list, description="Prevent assignment on Monday after full weekends") + prevent_monday_and_tuesday_after_full_weekends: List[int] = Field(default_factory=list, description="Prevent Monday/Tuesday after full weekends") + prevent_fridays_before_full_weekends: List[int] = Field(default_factory=list, description="Prevent Fridays before full weekends") + prevent_thursdays_before_full_weekends: List[int] = Field(default_factory=list, description="Prevent Thursdays before full weekends") + hard_constrain_pair_separation: bool = Field(False, description="Enforce hard pair separation constraint") + # These are lists of dicts with keys like {"grades": [...], "weeks": [...], "shifts": [...]} + avoid_shifts_by_grades: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning grades to shifts") + avoid_shifts_by_worker_names: List[dict] = Field(default_factory=list, description="List of rules to avoid assigning specific workers to shifts") + allocate_locum_shifts: bool = Field(True, description="Allow allocation of locum shifts") + distribute_locum_shifts: bool = Field(True, description="Distribute locum shifts among available workers") + balance_locum_shifts: bool = Field(False, description="Balance locum shifts similarly to normal shifts") + maximum_allowed_locum_shifts_per_worker: Optional[int] = Field(None, description="Maximum locum shifts allowed per worker (None = no limit)") + minimum_allowed_locum_shifts_per_worker: Optional[int] = Field(None, description="Minimum locum shifts required per worker (None = no minimum)") + min_summed_grade_by_shifts_per_day: List[MinSummedGradeByShiftsPerDayConstraint] = Field(default_factory=list, description="List of minimum-summed-grade constraints per day") + + # (Removed legacy dict-like convenience methods to enforce + # attribute access on the typed `RotaConstraintOptions` model.) class RotaBuilder(object): """Class to hold and manipulate shifts""" @@ -412,45 +414,21 @@ class RotaBuilder(object): self.balance_offset_modifier = balance_offset_modifier self.ltft_balance_offset = ltft_balance_offset - self.constraint_options = { - "balance_nights": True, - "constrain_time_off_after_nights": False, - "balance_nights_across_sites": True, - "balance_bank_holidays": True, - "balance_blocks": True, - "balance_shifts": True, # Does not use a quadratic function - "balance_shifts_true_quadratic": False, # Does not use a quadratic function - "balance_shifts_quadratic": False, # Will prevent spreading of spreading across different shifts - "balance_shifts_over_workers": True, - "minimise_shift_diffs": False, # less sophisticated version of balance_shifts_over_workers - "maximum_allowed_shift_diff": None, - "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": 1, - "max_night_frequency_week_exclusions": [], - "max_shifts_per_week": 7, - "max_shifts_per_month": 40, - "max_days_per_week_block": [], - #"max_days_per_week_block": [(3, 2)], # (max_days, week_block) - # The following will cause an unsolvable problem if a shift - # is forced as a block and spans the time peroid - "prevent_monday_after_full_weekends": [], - "prevent_monday_and_tuesday_after_full_weekends": [], - "prevent_fridays_before_full_weekends": [], - "prevent_thursdays_before_full_weekends": [], - "hard_constrain_pair_separation": False, - "avoid_shifts_by_grades": [], - "avoid_shifts_by_worker_names": [], - "allocate_locum_shifts": True, - "distribute_locum_shifts": True, - "balance_locum_shifts": False, - "maximum_allowed_locum_shifts_per_worker": None, - "minimum_allowed_locum_shifts_per_worker": None, - } - self.constraint_options["min_summed_grade_by_shifts_per_day"]: List[MinSummedGradeByShiftsPerDayConstraint] = [] + # Initialize constraint options with typed Pydantic model. Accept either + # None (use defaults), a dict (raw values), or an existing model instance. + # Store the model instance in `self.constraint_options_model` and also + # expose it as `self.constraint_options` for minimal code changes. + constraint_options = None + if constraint_options is None: + opts_model = RotaConstraintOptions() + elif isinstance(constraint_options, RotaConstraintOptions): + opts_model = constraint_options + elif isinstance(constraint_options, dict): + opts_model = RotaConstraintOptions.model_validate(constraint_options) + else: + opts_model = RotaConstraintOptions() + + self.constraint_options_model: RotaConstraintOptions = opts_model self.terminate_on_warning = [ "Worker/duplicate id", @@ -486,10 +464,11 @@ class RotaBuilder(object): self.MAX_SHIFTS_PER_DAY = 2 def set_rota_constraint(self, constraint: str, value: Any): - if constraint not in self.constraint_options: + # Ensure the model has the attribute and set it on the typed model + if not hasattr(self.constraint_options_model, constraint): raise ValueError(f"Constraint {constraint} not valid") - self.constraint_options[constraint] = value + setattr(self.constraint_options_model, constraint, value) def set_rota_dates(self, start_date: datetime.date, weeks_to_rota: int): self.weeks_to_rota = weeks_to_rota @@ -738,7 +717,7 @@ class RotaBuilder(object): initialize=0, ) - if self.constraint_options["balance_bank_holidays"]: + if self.constraint_options_model.balance_bank_holidays: # Bank holidays self.model.bank_holiday_count = Var( ((worker.id) for worker in self.workers), @@ -769,7 +748,7 @@ class RotaBuilder(object): # Used to limit number of workers on night shift per site # if we force a binary it will in effect hard constrain to < 2 # from a single site - if self.constraint_options["balance_nights_across_sites"]: + if self.constraint_options_model.balance_nights_across_sites: self.model.night_per_site = Var( ( (week, shift.name, site) @@ -864,7 +843,7 @@ class RotaBuilder(object): initialize=0, ) - if self.constraint_options["balance_shifts_over_workers"]: + if self.constraint_options_model.balance_shifts_over_workers: self.model.worker_shift_count_t1 = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, @@ -878,8 +857,8 @@ class RotaBuilder(object): ) if ( - self.constraint_options["balance_shifts"] - or self.constraint_options["balance_shifts_quadratic"] + self.constraint_options_model.balance_shifts + or self.constraint_options_model.balance_shifts_quadratic ): self.model.shift_count_t1 = Var( ( @@ -901,7 +880,7 @@ class RotaBuilder(object): initialize=0, ) - if self.constraint_options["balance_shifts_quadratic"]: + if self.constraint_options_model.balance_shifts_quadratic: self.model.shift_count_w = Var( ( (worker.id, shift.name) @@ -912,7 +891,7 @@ class RotaBuilder(object): initialize=0, ) - if self.constraint_options["balance_nights"]: + if self.constraint_options_model.balance_nights: # We also try to even out the night shifts seperately self.model.night_shift_count = Var( ((worker.id) for worker in self.workers), @@ -985,7 +964,7 @@ class RotaBuilder(object): bounds=self.SHIFT_BOUNDS["weekend_count"], ) - if self.constraint_options["balance_weekends"]: + if self.constraint_options_model.balance_weekends: self.model.weekend_shift_count_t1 = Var( ((worker.id) for worker in self.workers), domain=NonNegativeReals, @@ -1193,7 +1172,7 @@ class RotaBuilder(object): self.model.constraints = ConstraintList() # Create a set of constraints - for constraint in self.constraint_options.get("min_summed_grade_by_shifts_per_day", []): + for constraint in self.constraint_options_model.min_summed_grade_by_shifts_per_day: weeks = constraint.weeks if constraint.weeks is not None else self.weeks days_list = constraint.days if constraint.days is not None else self.days for week in weeks: @@ -1608,7 +1587,7 @@ class RotaBuilder(object): # Use per-worker value if set, otherwise use global constraint option if set max_days_worked_per_week = getattr(worker, "max_days_worked_per_week", None) if max_days_worked_per_week is None: - max_days_worked_per_week = self.constraint_options.get("max_days_worked_per_week", None) + max_days_worked_per_week = getattr(self.constraint_options_model, "max_days_worked_per_week", None) if max_days_worked_per_week is not None: for week in self.weeks: # Create a binary variable for each day to indicate if the worker works any shift that day @@ -1834,7 +1813,7 @@ class RotaBuilder(object): ) ) - if self.constraint_options["distribute_locum_shifts"]: + if self.constraint_options_model.distribute_locum_shifts: self.model.constraints.add( sum( self.model.locum_required[week, day, shift] @@ -1845,32 +1824,26 @@ class RotaBuilder(object): ) if ( - self.constraint_options["maximum_allowed_locum_shifts_per_worker"] + self.constraint_options_model.maximum_allowed_locum_shifts_per_worker is not None ): self.model.constraints.add( - self.constraint_options[ - "maximum_allowed_locum_shifts_per_worker" - ] + self.constraint_options_model.maximum_allowed_locum_shifts_per_worker >= self.model.locum_shifts_by_worker[worker.id] ) if worker.locum_availability: if ( - self.constraint_options[ - "minimum_allowed_locum_shifts_per_worker" - ] + self.constraint_options_model.minimum_allowed_locum_shifts_per_worker is not None ): self.model.constraints.add( - self.constraint_options[ - "minimum_allowed_locum_shifts_per_worker" - ] + self.constraint_options_model.minimum_allowed_locum_shifts_per_worker <= self.model.locum_shifts_by_worker[worker.id] ) if ( - self.constraint_options["minimise_shift_diffs"] - or self.constraint_options["maximum_allowed_shift_diff"] is not None + self.constraint_options_model.minimise_shift_diffs + or self.constraint_options_model.maximum_allowed_shift_diff is not None ): self.model.constraints.add( self.model.shift_count_diff_summed[worker.id] @@ -1880,17 +1853,17 @@ class RotaBuilder(object): ) ) - if self.constraint_options["maximum_allowed_shift_diff"] is not None: + if self.constraint_options_model.maximum_allowed_shift_diff is not None: self.model.constraints.add( self.model.shift_count_diff_summed[worker.id] - <= self.constraint_options["maximum_allowed_shift_diff"] + <= self.constraint_options_model.maximum_allowed_shift_diff ) self.model.constraints.add( self.model.shift_count_diff_summed[worker.id] - >= -self.constraint_options["maximum_allowed_shift_diff"] + >= -self.constraint_options_model.maximum_allowed_shift_diff ) - for max_days, weeks in self.constraint_options["max_days_per_week_block"]: + for max_days, weeks in self.constraint_options_model.max_days_per_week_block: for week_blocks in self.get_week_block_iterator(weeks): self.model.constraints.add( sum( @@ -1905,7 +1878,7 @@ class RotaBuilder(object): # Prevent more than n number shifts per 4 weeks try: self.model.constraints.add( - self.constraint_options["max_shifts_per_month"] + self.constraint_options_model.max_shifts_per_month >= sum( self.model.works[worker.id, week, day, shiftname] for week, day, shiftname in self.get_all_shiftname_combinations() @@ -1921,9 +1894,7 @@ class RotaBuilder(object): ) pass - for constraint_dict in self.constraint_options[ - "avoid_shifts_by_worker_names" - ]: + for constraint_dict in self.constraint_options_model.avoid_shifts_by_worker_names: if invalid_shifts := set(constraint_dict["shifts"]).difference( set(self.get_shift_names()) ): @@ -1945,7 +1916,7 @@ class RotaBuilder(object): ) ) - for constraint_dict in self.constraint_options["avoid_shifts_by_grades"]: + for constraint_dict in self.constraint_options_model.avoid_shifts_by_grades: if invalid_shifts := set(constraint_dict["shifts"]).difference( set(self.get_shift_names()) ): @@ -1980,9 +1951,9 @@ class RotaBuilder(object): == sum(self.model.works_weekend[worker.id, week] for week in self.weeks) ) - if self.constraint_options["max_weekends"] > -1: + if self.constraint_options_model.max_weekends > -1: self.model.constraints.add( - self.constraint_options["max_weekends"] + self.constraint_options_model.max_weekends >= self.model.worker_weekend_count[worker.id] ) @@ -2105,14 +2076,14 @@ class RotaBuilder(object): - worker.shift_target_number[shift.name] ) - if self.constraint_options["balance_shifts"]: + if self.constraint_options_model.balance_shifts: self.model.constraints.add( self.model.shift_count_t1[worker.id, shift.name] - self.model.shift_count_t2[worker.id, shift.name] == self.model.shift_count_diff[worker.id, shift.name] ) - if self.constraint_options["balance_shifts_quadratic"]: + if self.constraint_options_model.balance_shifts_quadratic: # This may need to be updated xU = 25 xL = 1 @@ -2159,7 +2130,7 @@ class RotaBuilder(object): # NOTE: as this balances across the sum of the worker's shifts # they can cancel each other out if a high weighting is given to multiple shifts - if self.constraint_options["balance_shifts_over_workers"]: + if self.constraint_options_model.balance_shifts_over_workers: self.model.constraints.add( self.model.worker_shift_count_t1[worker.id] - self.model.worker_shift_count_t2[worker.id] @@ -2170,7 +2141,7 @@ class RotaBuilder(object): ) ) - if self.constraint_options["balance_bank_holidays"]: + if self.constraint_options_model.balance_bank_holidays: extra_bank_holiday = 0 if self.use_bank_holiday_extra: extra_bank_holiday = worker.bank_holiday_extra @@ -2205,7 +2176,7 @@ class RotaBuilder(object): >= xU * (self.model.bank_holiday_count[worker.id] + 1) * 2 - xU * xU ) - if self.constraint_options["balance_nights"]: + if self.constraint_options_model.balance_nights: self.model.constraints.add( self.model.night_shift_count[worker.id] == sum( @@ -2280,7 +2251,7 @@ class RotaBuilder(object): worker.weekend_shift_target_number = weekend_shift_target_number - if self.constraint_options["balance_weekends"]: + if self.constraint_options_model.balance_weekends: if weekend_shift_target_number > 0: self.model.constraints.add( self.model.weekend_shift_count_t1[worker.id] @@ -2289,7 +2260,7 @@ class RotaBuilder(object): - weekend_shift_target_number ) - xU = self.constraint_options["max_weekends"] + xU = self.constraint_options_model.max_weekends xL = 1 self.model.constraints.add( @@ -2357,17 +2328,15 @@ class RotaBuilder(object): ] ) - if self.constraint_options["balance_blocks"]: - if self.constraint_options["max_night_frequency"]: + if self.constraint_options_model.balance_blocks: + if self.constraint_options_model.max_night_frequency: for week_blocks in self.get_week_block_iterator( - self.constraint_options["max_night_frequency"] + self.constraint_options_model.max_night_frequency ): # Ignore excluded weeks if set(week_blocks).intersection( set( - self.constraint_options[ - "max_night_frequency_week_exclusions" - ] + self.constraint_options_model.max_night_frequency_week_exclusions ) ): continue @@ -2387,10 +2356,10 @@ class RotaBuilder(object): ) ) - # if self.constraint_options["balance_weekends"]: - if self.constraint_options["max_weekend_frequency"]: + # if self.constraint_options_model.balance_weekends: + if self.constraint_options_model.max_weekend_frequency: for week_blocks in self.get_week_block_iterator( - self.constraint_options["max_weekend_frequency"] + self.constraint_options_model.max_weekend_frequency ): # Prevent weekend shifts more than once every n weeks self.model.constraints.add( @@ -2474,7 +2443,7 @@ class RotaBuilder(object): try: self.model.constraints.add( - self.constraint_options["max_shifts_per_week"] + self.constraint_options_model.max_shifts_per_week >= sum( self.model.works[worker.id, w, day, shiftname] # for shiftname in self.get_shift_names_by_week_day(week, day) @@ -2490,7 +2459,7 @@ class RotaBuilder(object): # TODO: consider excluding shifts that span relevant period if ( worker.site - in self.constraint_options["prevent_monday_after_full_weekends"] + in self.constraint_options_model.prevent_monday_after_full_weekends ): # Ignore last week if week + 1 in self.weeks: @@ -2517,9 +2486,7 @@ class RotaBuilder(object): ) if ( worker.site - in self.constraint_options[ - "prevent_monday_and_tuesday_after_full_weekends" - ] + in self.constraint_options_model.prevent_monday_and_tuesday_after_full_weekends ): # Ignore last week if week + 1 in self.weeks: @@ -2552,7 +2519,7 @@ class RotaBuilder(object): ) if ( worker.site - in self.constraint_options["prevent_fridays_before_full_weekends"] + in self.constraint_options_model.prevent_fridays_before_full_weekends ): self.model.constraints.add( full_weekend_count @@ -2571,7 +2538,7 @@ class RotaBuilder(object): ) if ( worker.site - in self.constraint_options["prevent_thursdays_before_full_weekends"] + in self.constraint_options_model.prevent_thursdays_before_full_weekends ): self.model.constraints.add( full_weekend_count @@ -2850,7 +2817,7 @@ class RotaBuilder(object): # IF paired we check the following against both workers workers = [worker] - if self.constraint_options["hard_constrain_pair_separation"]: + if self.constraint_options_model.hard_constrain_pair_separation: for worker_pairs in self.worker_pairs: if worker_pairs[0] == worker: workers = worker_pairs @@ -2958,7 +2925,7 @@ class RotaBuilder(object): ) ) - if self.constraint_options["hard_constrain_pair_separation"]: + if self.constraint_options_model.hard_constrain_pair_separation: for worker_pairs in self.worker_pairs: if worker_pairs[0] == worker: self.model.constraints.add( @@ -3130,7 +3097,7 @@ class RotaBuilder(object): locum_shift_balancing = 0 if self.get_workers_who_require_locums(): - if self.constraint_options["balance_locum_shifts"]: + if self.constraint_options_model.balance_locum_shifts: locum_shift_balancing = sum( locum_shift_balance_modifier_constant * ( @@ -3140,7 +3107,7 @@ class RotaBuilder(object): for worker in self.workers ) - if self.constraint_options["balance_shifts_over_workers"]: + if self.constraint_options_model.balance_shifts_over_workers: worker_shift_balancing = sum( balance_modifier_constant * ( @@ -3152,7 +3119,7 @@ class RotaBuilder(object): else: worker_shift_balancing = 0 - if self.constraint_options["balance_shifts"]: + if self.constraint_options_model.balance_shifts: shift_balancing = sum( balance_modifier_constant * ( @@ -3165,7 +3132,7 @@ class RotaBuilder(object): else: shift_balancing = 0 - if self.constraint_options["balance_shifts_quadratic"]: + if self.constraint_options_model.balance_shifts_quadratic: quadratic_shift_balancing = sum( balance_quadratic_shift_modifier_constant * ( @@ -3179,7 +3146,7 @@ class RotaBuilder(object): quadratic_shift_balancing = 0 true_quadratic_shift_balancing = 0 - if self.constraint_options["balance_shifts_true_quadratic"]: + if self.constraint_options_model.balance_shifts_true_quadratic: shift_diff_modifier_constant = 10 true_quadratic_shift_balancing = sum( shift_diff_modifier_constant @@ -3187,7 +3154,7 @@ class RotaBuilder(object): for worker in self.workers ) - if self.constraint_options["minimise_shift_diffs"]: + if self.constraint_options_model.minimise_shift_diffs: shift_diff_modifier_constant = 10 shift_diff_balancing = sum( shift_diff_modifier_constant @@ -3197,7 +3164,7 @@ class RotaBuilder(object): else: shift_diff_balancing = 0 - if self.constraint_options["balance_nights"]: + if self.constraint_options_model.balance_nights: night_balance_modifier_constant = 10 night_shift_balancing = sum( night_balance_modifier_constant @@ -3209,7 +3176,7 @@ class RotaBuilder(object): else: night_shift_balancing = 0 - if self.constraint_options["balance_bank_holidays"]: + if self.constraint_options_model.balance_bank_holidays: bank_holiday_balance_modifier_constant = 1000 bank_holiday_balancing = sum( bank_holiday_balance_modifier_constant @@ -3221,7 +3188,7 @@ class RotaBuilder(object): else: bank_holiday_balancing = 0 - if self.constraint_options["balance_weekends"]: + if self.constraint_options_model.balance_weekends: weekend_balance_modifier_constant = 5 weekend_shift_balancing = sum( weekend_balance_modifier_constant @@ -3252,7 +3219,7 @@ class RotaBuilder(object): ) # # Spread nights - if self.constraint_options["balance_nights_across_sites"]: + if self.constraint_options_model.balance_nights_across_sites: nights_site_balancing = sum( ( self.model.night_per_site_t1[week, shift.name, site] @@ -3282,7 +3249,7 @@ class RotaBuilder(object): else: block_site_balancing = 0 - if self.constraint_options["balance_blocks"]: + if self.constraint_options_model.balance_blocks: blocks_balancing = sum( block_shift_balancing_constant * self.model.blocks_assigned[week, shift] @@ -3446,14 +3413,14 @@ class RotaBuilder(object): def add_worker_name_constraint_by_week( self, names: Iterable[str], weeks: Iterable[int], shifts ): - self.constraint_options["avoid_shifts_by_worker_names"].append( + self.constraint_options_model.avoid_shifts_by_worker_names.append( {"names": names, "weeks": weeks, "shifts": shifts} ) def add_grade_constraint_by_week( self, grades: Iterable[int], weeks: Iterable[int], shifts ): - self.constraint_options["avoid_shifts_by_grades"].append( + self.constraint_options_model.avoid_shifts_by_grades.append( {"grades": grades, "weeks": weeks, "shifts": shifts} ) @@ -3473,7 +3440,7 @@ class RotaBuilder(object): weeks=weeks, days=days, ) - self.constraint_options["min_summed_grade_by_shifts_per_day"].append(constraint) + self.constraint_options_model.min_summed_grade_by_shifts_per_day.append(constraint) def add_shift(self, shift: SingleShift) -> None: """Add a shift to the collection @@ -4558,7 +4525,7 @@ class RotaBuilder(object): + f"#weekends_worked: {model.worker_weekend_count[worker.id].value}\\#" ) - if self.constraint_options["balance_bank_holidays"]: + if self.constraint_options_model.balance_bank_holidays: bank_holiday_count = model.bank_holiday_count[worker.id].value bank_holiday_count_w = model.bank_holiday_count_w[worker.id].value - 1