diff --git a/rota/shifts.py b/rota/shifts.py index 619686f..3449332 100644 --- a/rota/shifts.py +++ b/rota/shifts.py @@ -143,6 +143,7 @@ class SingleShift(BaseModel): start_date: datetime.date | None = None end_date: datetime.date | None = None pair_proxy: bool = False + display_char: str | None = None # Add this line model_config = ConfigDict( extra="allow", @@ -160,6 +161,9 @@ class SingleShift(BaseModel): # versus target shifts (hard constraint) # if 0 exactly equal shifts must be assigend (unlikely to be possible) + if self.display_char is not None and len(self.display_char) > 1: + raise ValueError("display_char must be a single character or None") + if self.balance_offset is None: self.balance_offset = len(self.days) @@ -215,6 +219,12 @@ class SingleShift(BaseModel): return 0 + def get_display_char(self): + """Returns the display character for the shift""" + if self.display_char is not None: + return self.display_char + return self.name[0] + class RotaBuilder(object): """Class to hold and manipulate shifts""" @@ -873,8 +883,6 @@ class RotaBuilder(object): def locum_request_init(model, wid, week, day, shift): - print("locum rquest init") - print(locals()) if (wid, week, day, shift) in self.locum_availability: self.locum_availability_map[(wid, week, day)] = shift return 1 @@ -1896,14 +1904,16 @@ class RotaBuilder(object): worker.id, week, day, shift.name ] ) - if not worker.locum_on_nwds and day == n: - if start <= self.week_day_date_map[(week, day)] < end: - self.model.constraints.add( - 0 - == self.model.locum_works[ - worker.id, week, day, shift.name - ] - ) + + if self.get_workers_who_require_locums(): + if not worker.locum_on_nwds and day == n: + if start <= self.week_day_date_map[(week, day)] < end: + self.model.constraints.add( + 0 + == self.model.locum_works[ + worker.id, week, day, shift.name + ] + ) if self.constraint_options["balance_blocks"]: if self.constraint_options["max_night_frequency"]: @@ -3497,9 +3507,10 @@ class RotaBuilder(object): for week, day in self.get_week_day_combinations(): a = "-" for shift in self.get_shift_names_by_week_day(week, day): + shift_obj = self.get_shift_by_name(shift) if model.works[worker.id, week, day, shift].value > 0: shifts.append(shift) - a = shift[0] + a = shift_obj.get_display_char() w.append(a) shift_count = "" @@ -3616,15 +3627,16 @@ class RotaBuilder(object): # if the worker has been assigned locum = False for shift in self.get_shift_names_by_week_day(week, day): + shift_obj = self.get_shift_by_name(shift) if model.works[worker.id, week, day, shift].value > 0.8: shifts.append(shift) - a = shift[0] + a = shift_obj.get_display_char() shift_name = shift break if self.get_workers_who_require_locums(): if model.locum_works[worker.id, week, day, shift].value > 0.8: locum_shifts.append(shift) - a = shift[0] + a = shift_obj.get_display_char() shift_name = shift locum = True @@ -3872,9 +3884,9 @@ class RotaBuilder(object): } for worker in self.workers: for shift in self.get_shift_names(): - timetable[worker.name][shift] = self.model.shift_count[ + timetable[worker.name][shift] = round(self.model.shift_count[ worker.id, shift - ].value + ].value) return timetable @@ -3915,7 +3927,7 @@ class RotaBuilder(object): for shift in self.get_shift_names_by_week_day(week, day): if self.model.works[worker.id, week, day, shift].value > 0.90: temp = shift - if include_locums and self.model.locum_works[worker.id, week, day, shift].value > 0.90: + elif include_locums and self.model.locum_works[worker.id, week, day, shift].value > 0.90: temp = shift shifts.append(temp) @@ -3925,7 +3937,7 @@ class RotaBuilder(object): shifts = self.get_worker_shift_list(worker) # Convert shift to a string representation - return "".join([i if i != "" else "-" for i in shifts]) + return "".join([self.get_shift_by_name(i).get_display_char() if i != "" else "-" for i in shifts]) def get_shift_summary(self): works = self.model.works @@ -3957,7 +3969,7 @@ class RotaBuilder(object): def get_shift_summary_html(self): works = self.model.works timetable = { - worker.get_details(): {shift[0]: "" for shift in self.get_shift_names()} + worker.get_details(): {shift.get_display_char(): "" for shift in self.get_shifts()} for worker in self.workers } # timetable = { worker.get_details() : { shift : "" for shift in ["truro_twilight"] } for worker in workers } @@ -3977,7 +3989,8 @@ class RotaBuilder(object): l.append(f"{shift} ({c})") total_shifts = total_shifts + c # print(worker.id, shift) - timetable[worker.get_details()][shift[0]] = c + shift_obj = self.get_shift_by_name(shift) + timetable[worker.get_details()][shift_obj.get_display_char()] = c t.append(f"{worker.get_full_details()} [{total_shifts}]: {', '.join(l)}") return "\n".join(t) diff --git a/test/test_shifts.py b/test/test_shifts.py index 2bec64a..8dc18b9 100644 --- a/test/test_shifts.py +++ b/test/test_shifts.py @@ -708,4 +708,25 @@ class TestShiftWorkerRequirements: total_shifts = Rota.get_workers_total_shifts() for worker in Rota.get_workers(): - assert total_shifts[worker.name] == 30 \ No newline at end of file + assert total_shifts[worker.name] == 30 + +class TestDisplayChar: + def test_display_char_in_output(monkeypatch): + Rota = generate_basic_rota(workers=1) + shift = SingleShift( + sites=("group1",), + name="testshift", + length=8, + days=days[:1], + display_char="Z" + ) + Rota.add_shifts(shift) + + assert shift.display_char == "Z" + + Rota.build_and_solve(options={"ratio": 0.1}) + + assert Rota.results.solver.status == "ok" + + for worker in Rota.get_workers(): + assert Rota.get_worker_shift_list_string(worker).startswith("Z------" * 5) \ No newline at end of file