add custom display char for shift

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,6 +1904,8 @@ class RotaBuilder(object):
                                         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(
@@ -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)
 

test/test_shifts.py

@@ -709,3 +709,24 @@ class TestShiftWorkerRequirements:
 
         for worker in Rota.get_workers():
             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)