add prtial shifts

rota_generator/workers.py

@@ -169,6 +169,42 @@ class MaxUniqueShiftsPerWeekBlockConstraint(BaseModel):
     )
 
 
+class ShiftStartDate(BaseModel):
+    shift: str
+    start_date: datetime.date
+
+    @field_validator("start_date", mode="before")
+    @classmethod
+    def coerce_date(cls, v):
+        if isinstance(v, datetime.date):
+            return v
+        if isinstance(v, str):
+            for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"):
+                try:
+                    return datetime.datetime.strptime(v, fmt).date()
+                except Exception:
+                    continue
+        raise ValueError(f"Cannot parse date: {v}")
+
+
+class ShiftEndDate(BaseModel):
+    shift: str
+    end_date: datetime.date
+
+    @field_validator("end_date", mode="before")
+    @classmethod
+    def coerce_date(cls, v):
+        if isinstance(v, datetime.date):
+            return v
+        if isinstance(v, str):
+            for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"):
+                try:
+                    return datetime.datetime.strptime(v, fmt).date()
+                except Exception:
+                    continue
+        raise ValueError(f"Cannot parse date: {v}")
+
+
 class Worker(BaseModel):
     name: str
     site: str
@@ -215,11 +251,55 @@ class Worker(BaseModel):
 
     exact_shifts: dict[str, int] = {} # Map shift_name to exact number of shifts
     groups: set[str] = set() # Groups that the worker belongs to
+    shift_start_dates: list[ShiftStartDate] = [] # List of ShiftStartDate models
+    shift_end_dates: list[ShiftEndDate] = [] # List of ShiftEndDate models
 
     weekend_shift_target_number: float = 0.0
     avoid_shifts_on_dates: list[AvoidShiftOnDates] = []
 
+    @field_validator("shift_start_dates", mode="before")
+    @classmethod
+    def coerce_shift_start_dates(cls, v):
+        if v is None:
+            return []
+        if isinstance(v, dict):
+            res = []
+            for shift, d in v.items():
+                res.append(ShiftStartDate(shift=shift, start_date=d))
+            return res
+        if isinstance(v, list):
+            res = []
+            for item in v:
+                if isinstance(item, ShiftStartDate):
+                    res.append(item)
+                elif isinstance(item, dict):
+                    res.append(ShiftStartDate(**item))
+                else:
+                    raise ValueError(f"Invalid ShiftStartDate item: {item}")
+            return res
+        raise ValueError("Must be a list or dictionary")
 
+    @field_validator("shift_end_dates", mode="before")
+    @classmethod
+    def coerce_shift_end_dates(cls, v):
+        if v is None:
+            return []
+        if isinstance(v, dict):
+            res = []
+            for shift, d in v.items():
+                res.append(ShiftEndDate(shift=shift, end_date=d))
+            return res
+        if isinstance(v, list):
+            res = []
+            for item in v:
+                if isinstance(item, ShiftEndDate):
+                    res.append(item)
+                elif isinstance(item, dict):
+                    res.append(ShiftEndDate(**item))
+                else:
+                    raise ValueError(f"Invalid ShiftEndDate item: {item}")
+            return res
+        raise ValueError("Must be a list or dictionary")
 
     model_config = ConfigDict(
         extra="allow",
@@ -416,13 +496,88 @@ class Worker(BaseModel):
         self.proportion_rota_to_work = days_to_work / Rota.rota_days_length
         self.days_to_work = days_to_work
 
+        # Shift-specific start/end dates and adjusted FTEs
+        self.calculated_shift_start_dates = {}
+        self.calculated_shift_end_dates = {}
+        self.proportion_rota_to_work_shifts = {}
+
+        # Convert list of ShiftStartDate/ShiftEndDate models/dicts to dictionary mapping for quick lookup
+        start_dates_dict = {item.shift: item.start_date for item in getattr(self, "shift_start_dates", []) if hasattr(item, "shift")}
+        end_dates_dict = {item.shift: item.end_date for item in getattr(self, "shift_end_dates", []) if hasattr(item, "shift")}
+
+        for shift in Rota.get_shifts():
+            s_date = start_dates_dict.get(shift.name, self.start_date)
+            if s_date is None:
+                calc_s = self.calculated_start_date
+            else:
+                if isinstance(s_date, str):
+                    for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"):
+                        try:
+                            s_date = datetime.datetime.strptime(s_date, fmt).date()
+                            break
+                        except Exception:
+                            continue
+                calc_s = s_date
+                if calc_s < Rota.start_date:
+                    calc_s = Rota.start_date
+                elif calc_s > Rota.rota_end_date:
+                    calc_s = Rota.rota_end_date
+
+            e_date = end_dates_dict.get(shift.name, self.end_date)
+            if e_date is None:
+                calc_e = self.calculated_end_date
+            else:
+                if isinstance(e_date, str):
+                    for fmt in ("%d/%m/%Y", "%d/%m/%y", "%Y-%m-%d"):
+                        try:
+                            e_date = datetime.datetime.strptime(e_date, fmt).date()
+                            break
+                        except Exception:
+                            continue
+                calc_e = e_date
+                if calc_e > Rota.rota_end_date:
+                    calc_e = Rota.rota_end_date
+                elif calc_e < Rota.start_date:
+                    calc_e = Rota.start_date
+
+            if calc_s >= calc_e:
+                calc_s = Rota.rota_end_date
+                calc_e = Rota.rota_end_date
+
+            self.calculated_shift_start_dates[shift.name] = calc_s
+            self.calculated_shift_end_dates[shift.name] = calc_e
+
+            days_active = (calc_e - calc_s).days
+
+            # Subtract overlapping OOP days from the active period of this shift
+            for item in self.oop:
+                start_oop = item.start_date
+                end_oop = item.end_date
+                if isinstance(start_oop, datetime.date):
+                    start_oop_date = start_oop
+                else:
+                    start_oop_date = datetime.datetime.strptime(start_oop, "%d/%m/%y").date()
+
+                if isinstance(end_oop, datetime.date):
+                    end_oop_date = end_oop
+                else:
+                    end_oop_date = datetime.datetime.strptime(end_oop, "%d/%m/%y").date()
+
+                overlap_start = max(calc_s, start_oop_date)
+                overlap_end = min(calc_e, end_oop_date)
+                if overlap_start < overlap_end:
+                    days_active -= (overlap_end - overlap_start).days
+
+            self.proportion_rota_to_work_shifts[shift.name] = max(0.0, days_active / Rota.rota_days_length)
+
         # We have to adjust the full time equivalent for people who CCT / leave the rota early
         self.fte_adj = self.fte * self.proportion_rota_to_work
 
         self.fte_adj_shifts = {}
-        if self.shift_fte_overrides:
-            for shift, fte in self.shift_fte_overrides.items():
-                self.fte_adj_shifts[shift] = fte * self.proportion_rota_to_work
+        for shift in Rota.get_shifts():
+            prop = self.proportion_rota_to_work_shifts.get(shift.name, self.proportion_rota_to_work)
+            base_fte = self.shift_fte_overrides.get(shift.name, self.fte)
+            self.fte_adj_shifts[shift.name] = base_fte * prop
 
 
         if self.fte_adj > 100: