ross/proc-rota
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: implement sub-block handling for shifts and update constraints in RotaBuilder

Browse Source
This commit is contained in:
Ross
2026-07-10 08:08:51 +01:00
parent d1fbbe2dc9
commit 5ff78faa07
3 changed files with 191 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gen_proc.py
+4 -3
View File
@@ -74,6 +74,7 @@ def main(
Rota.constraint_options["max_weekend_frequency"] = 3 Rota.constraint_options["max_weekend_frequency"] = 3
Rota.constraint_options["max_days_per_week_block"] = [(8,3)] Rota.constraint_options["max_days_per_week_block"] = [(8,3)]
Rota.constraint_options["maximum_allowed_shift_diff"] = 1
# wr = [ # wr = [
# WorkerRequirement( # WorkerRequirement(
@@ -174,11 +175,11 @@ def main(
length=12.5, length=12.5,
days=days[4:], days=days[4:],
# balance_offset=3, # balance_offset=3,
rota_on_nwds=True, #rota_on_nwds=True,
force_as_block=True, #force_as_block=True,
# assign_as_block=True, # assign_as_block=True,
constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)], constraints=[PreShiftConstraint(days=2), PostShiftConstraint(days=2)],
# force_as_block_unless_nwd=True force_as_block_unless_nwd=[["Fri"], ["Sat", "Sun"]],
), ),
SingleShift( SingleShift(
sites=("torbay", "torbay twilights and weekends"), sites=("torbay", "torbay twilights and weekends"),
rota_generator/shifts.py
+128 -92
View File
@@ -417,7 +417,7 @@ class SingleShift(BaseModel):
rota_on_nwds: bool = False rota_on_nwds: bool = False
assign_as_block: bool = False assign_as_block: bool = False
force_as_block: bool = False force_as_block: bool = False
force_as_block_unless_nwd: bool = False force_as_block_unless_nwd: bool | List[List[str]] = False
hard_constrain_shift: bool = True hard_constrain_shift: bool = True
bank_holidays_only: bool = False bank_holidays_only: bool = False
#constraint: list[ShiftConstraint] = [] #constraint: list[ShiftConstraint] = []
@@ -1050,20 +1050,28 @@ class RotaBuilder(object):
initialize=0, initialize=0,
) )
blocks_worker_keys = []
for worker, week, shift_name in self.worker_week_shifts_to_assign_as_blocks():
shift = self.get_shift_by_name(shift_name)
sub_blocks = self.get_shift_sub_blocks(shift)
for sb_idx in range(len(sub_blocks)):
blocks_worker_keys.append((worker.id, week, shift_name, sb_idx))
self.model.blocks_worker_shift_assigned = Var( self.model.blocks_worker_shift_assigned = Var(
( blocks_worker_keys,
(worker.id, week, shift)
for worker, week, shift in self.worker_week_shifts_to_assign_as_blocks()
),
within=Binary, within=Binary,
initialize=0, initialize=0,
) )
blocks_assigned_keys = []
for week, shift_name in self.week_shifts_to_assign_as_blocks():
shift = self.get_shift_by_name(shift_name)
sub_blocks = self.get_shift_sub_blocks(shift)
for sb_idx in range(len(sub_blocks)):
blocks_assigned_keys.append((week, shift_name, sb_idx))
self.model.blocks_assigned = Var( self.model.blocks_assigned = Var(
( blocks_assigned_keys,
(week, shift)
for week, shift in self.week_shifts_to_assign_as_blocks()
),
within=NonNegativeIntegers, within=NonNegativeIntegers,
initialize=0, initialize=0,
) )
@@ -1900,99 +1908,103 @@ class RotaBuilder(object):
for week, shift_name in self.week_shifts_to_assign_as_blocks(): for week, shift_name in self.week_shifts_to_assign_as_blocks():
shift = self.get_shift_by_name(shift_name) shift = self.get_shift_by_name(shift_name)
sub_blocks = self.get_shift_sub_blocks(shift)
for worker in self.get_workers_for_shift(shift): for sb_idx, sb in enumerate(sub_blocks):
if (worker.id, week, shift.name) in self.model.blocks_worker_shift_assigned: for worker in self.get_workers_for_shift(shift):
try: if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned:
self.model.constraints.add( try:
8 self.model.constraints.add(
* self.model.blocks_worker_shift_assigned[ 8
worker.id, week, shift_name * self.model.blocks_worker_shift_assigned[
] worker.id, week, shift_name, sb_idx
>= sum( ]
self.model.works[worker.id, week, day, shift_name] >= sum(
for day in shift.days self.model.works[worker.id, week, day, shift_name]
for day in sb
)
) )
) except KeyError:
except KeyError: pass
pass
self.model.constraints.add( self.model.constraints.add(
self.model.blocks_assigned[week, shift_name] self.model.blocks_assigned[week, shift_name, sb_idx]
== sum( == sum(
self.model.blocks_worker_shift_assigned[ self.model.blocks_worker_shift_assigned[
worker.id, week, shift_name worker.id, week, shift_name, sb_idx
] ]
for worker in self.workers for worker in self.workers
if (worker.id, week, shift.name) in self.model.blocks_worker_shift_assigned if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
)
) )
)
workers_required = shift.get_worker_requirement_by_date( workers_required = shift.get_worker_requirement_by_date(
self.get_week_start_date(week) self.get_week_start_date(week)
) )
if shift.force_as_block_unless_nwd: if shift.force_as_block_unless_nwd:
workers = self.get_workers_for_shift(shift) workers = self.get_workers_for_shift(shift)
# Get workers who have a nwd on the shift # Get workers who have a nwd on the sub-block
nwd_workers = [] nwd_workers = []
full_workers = [] full_workers = []
for w in workers: for w in workers:
if w.non_working_day_list: if w.non_working_day_list:
l = [] l = []
for ( for (
nwd, nwd,
start_nwd_date, start_nwd_date,
end_nwd_date, end_nwd_date,
) in w.non_working_day_list: ) in w.non_working_day_list:
if nwd in shift.days: if nwd in sb:
if start_nwd_date > self.get_week_start_date(week): if start_nwd_date > self.get_week_start_date(week):
continue continue
if end_nwd_date < self.get_week_start_date(week): if end_nwd_date < self.get_week_start_date(week):
continue continue
l.append(w) l.append(w)
if l: if l:
nwd_workers.append(w) nwd_workers.append(w)
else:
full_workers.append(w)
else: else:
full_workers.append(w) full_workers.append(w)
if not nwd_workers: # Just do the usual
self.model.constraints.add(
sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift.name, sb_idx
]
for worker in self.workers
if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
)
<= workers_required
)
else: else:
full_workers.append(w) self.model.constraints.add(
sum(
if not nwd_workers: # Just do the usual self.model.blocks_worker_shift_assigned[
self.model.constraints.add( worker.id, week, shift.name, sb_idx
sum( ]
self.model.blocks_worker_shift_assigned[ for worker in full_workers
worker.id, week, shift.name if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
] )
for worker in self.workers <= workers_required + 1
) )
<= workers_required
)
else: elif shift.force_as_block:
self.model.constraints.add( if workers_required:
sum( self.model.constraints.add(
self.model.blocks_worker_shift_assigned[ sum(
worker.id, week, shift.name self.model.blocks_worker_shift_assigned[
] worker.id, week, shift.name, sb_idx
for worker in full_workers ]
for worker in self.workers
if (worker.id, week, shift.name, sb_idx) in self.model.blocks_worker_shift_assigned
)
<= workers_required
) )
<= workers_required + 1
)
elif shift.force_as_block:
if workers_required:
self.model.constraints.add(
sum(
self.model.blocks_worker_shift_assigned[
worker.id, week, shift.name
]
for worker in self.workers
)
<= workers_required
)
# Most of our constraints apply per worker # Most of our constraints apply per worker
# Worker constraint loop (worker loop) # Worker constraint loop (worker loop)
@@ -4185,24 +4197,27 @@ class RotaBuilder(object):
if self.constraint_options_model.balance_blocks: if self.constraint_options_model.balance_blocks:
blocks_balancing = sum( blocks_balancing = sum(
block_shift_balancing_constant block_shift_balancing_constant
* self.model.blocks_assigned[week, shift] * self.model.blocks_assigned[week, shift, sb_idx]
for week in self.weeks for week in self.weeks
for shift in self.shifts_to_assign_as_blocks() for shift in self.shifts_to_assign_as_blocks()
for sb_idx in range(len(self.get_shift_sub_blocks(self.get_shift_by_name(shift))))
if (week, shift, sb_idx) in self.model.blocks_assigned
) )
else: else:
blocks_balancing = 0 blocks_balancing = 0
prefer_block_expr = sum( prefer_block_expr = sum(
worker.assign_as_block_preferences.get(shift.name, 0) worker.assign_as_block_preferences.get(shift.name, 0)
* self.model.blocks_worker_shift_assigned[worker.id, week, shift.name] * self.model.blocks_worker_shift_assigned[worker.id, week, shift.name, sb_idx]
for worker in self.workers for worker in self.workers
for week in self.weeks for week in self.weeks
for shift in self.shifts for shift in self.shifts
if hasattr(worker, "assign_as_block_preferences") if hasattr(worker, "assign_as_block_preferences")
and shift.name in worker.assign_as_block_preferences and shift.name in worker.assign_as_block_preferences
and (worker.id, week, shift.name)
in self.model.blocks_worker_shift_assigned
and self.is_worker_eligible_for_shift(worker, shift) and self.is_worker_eligible_for_shift(worker, shift)
for sb_idx in range(len(self.get_shift_sub_blocks(shift)))
if (worker.id, week, shift.name, sb_idx)
in self.model.blocks_worker_shift_assigned
) )
# Quadratic :( # Quadratic :(
@@ -5046,6 +5061,27 @@ class RotaBuilder(object):
) )
]) ])
def get_shift_sub_blocks(self, shift: SingleShift) -> List[List[str]]:
if getattr(shift, "force_as_block_unless_nwd", None):
val = shift.force_as_block_unless_nwd
if isinstance(val, (list, tuple)) and all(isinstance(x, (list, tuple, set)) for x in val):
return [list(x) for x in val]
return [list(shift.days)]
# Check if it is a block shift globally or per-worker
is_block = (
shift.assign_as_block
or shift.force_as_block
or any(
hasattr(worker, "assign_as_block_preferences")
and getattr(worker, "assign_as_block_preferences", {}).get(shift.name, 0) != 0
for worker in self.workers
)
)
if is_block:
return [list(shift.days)]
return []
def get_all_locum_availability(self): def get_all_locum_availability(self):
return self.locum_availability_map return self.locum_availability_map
test/test_nwds.py
+59 -2
View File
@@ -163,7 +163,7 @@ def test_nwd_force_as_block_force_split():
for worker in Rota.workers: for worker in Rota.workers:
shifts = Rota.get_worker_shift_list(worker) shifts = Rota.get_worker_shift_list(worker)
shifts_string = "".join([i if i != "" else "-" for i in shifts]) shifts_string = "".join([i if i != "" else "-" for i in shifts])
for week in weeks_from_list(shifts_string[7 * 5 :]): for week in weeks_from_list(shifts_string[7 * 6 :]):
assert week in ("-----ww", "dddddww") assert week in ("-----ww", "dddddww")
if worker.name == "worker1": if worker.name == "worker1":
assert shifts_string[: 7 * 5].count("ww-") == 4 assert shifts_string[: 7 * 5].count("ww-") == 4
@@ -213,4 +213,61 @@ def test_nwd_testing():
if worker.name == "worker1": if worker.name == "worker1":
assert week == "dddddww" assert week == "dddddww"
else: else:
assert week == "-----ww" assert week == "-----ww"
def test_force_as_block_unless_nwd_with_subblocks():
# Test that a Fri/Sat/Sun shift can be split into Fri and Sat/Sun blocks
# if a worker does not work on Fridays.
Rota = setup_rota(weeks_to_rota=2)
start_date = Rota.start_date
# worker1 does not work on Fridays (Fri is NWD)
worker1 = Worker(
name="worker1", site="group1", grade=1,
nwds=[{"day": "Fri", "start_date": start_date, "end_date": start_date + datetime.timedelta(weeks=2)}],
)
# worker2 works normally
worker2 = Worker(name="worker2", site="group1", grade=1)
Rota.add_workers((worker1, worker2))
# Fri/Sat/Sun shift, requiring 1 worker, forced as block unless NWD
# We specify sub-blocks: [["Fri"], ["Sat", "Sun"]]
Rota.add_shifts(
SingleShift(
sites=("group1",), name="weekend", length=12.5,
days=["Fri", "Sat", "Sun"],
workers_required=1,
force_as_block_unless_nwd=[["Fri"], ["Sat", "Sun"]],
),
)
# We remove no valid shifts warning
Rota.terminate_on_warning.remove("Worker/no valid shifts")
Rota.build_and_solve()
assert Rota.results.solver.status == "ok"
# Check assignments:
# Since worker1 has NWD on Friday:
# - worker1 cannot work Friday.
# - But worker1 can work Sat and Sun as a block!
# So on Saturday and Sunday, worker1 should be assigned.
# On Friday, worker2 should be assigned.
for week in Rota.weeks:
w1_fri = Rota.model.works[worker1.id, week, "Fri", "weekend"].value
w1_sat = Rota.model.works[worker1.id, week, "Sat", "weekend"].value
w1_sun = Rota.model.works[worker1.id, week, "Sun", "weekend"].value
w2_fri = Rota.model.works[worker2.id, week, "Fri", "weekend"].value
w2_sat = Rota.model.works[worker2.id, week, "Sat", "weekend"].value
w2_sun = Rota.model.works[worker2.id, week, "Sun", "weekend"].value
assert w1_fri == 0
# Either worker1 works Sat/Sun and worker2 works Fri (split block)
# OR worker2 works Fri/Sat/Sun (full block) and worker1 works nothing
is_split = (w1_sat == 1 and w1_sun == 1 and w2_fri == 1 and w2_sat == 0 and w2_sun == 0)
is_full_w2 = (w1_sat == 0 and w1_sun == 0 and w2_fri == 1 and w2_sat == 1 and w2_sun == 1)
assert is_split or is_full_w2