Refactor date parsing and configuration constants in gen_cons.py for improved clarity and maintainability

This commit is contained in:
Ross
2026-05-19 20:37:06 +01:00
parent 4bb4f5b97a
commit 68b80ace23
+59 -54
View File
@@ -22,16 +22,29 @@ from loguru import logger
app = typer.Typer() app = typer.Typer()
ROTA_REQUESTS = "cons/requests.ods"
ROTA_B_PATH = "cons/rota_b.xlsx"
ROTA_START_DATE = "2026-08-03"
sites = ("rota a", "rota b", "rota c") sites = ("rota a", "rota b", "rota c")
def _parse_sheet_date(value):
if pd.isna(value):
return None
parsed = pd.to_datetime(value, errors="coerce")
if pd.isna(parsed):
return None
return parsed.date()
def extract_leave_and_rota_from_calender(calender_df): def extract_leave_and_rota_from_calender(calender_df):
""" """
Extract leave and rota assignments from calender_df. Extract leave and rota assignments from calender_df.
- The first row contains worker names, starting from the 5th column (index 4). - The first row contains worker names, starting from the 5th column (index 4).
- The second column (index 1) contains dates in dd/mm/yyyy format. - The second column (index 1) contains dates in dd/mm/yyyy format.
- The second column (index 1) also contains the rota/leave code for that date. - The second column (index 1) also contains the rota/leave code for that date.
Returns a list of dicts: {"date": ..., "worker": ..., "rota": ...} Returns a list of dicts: {"date": ..., "/worker": ..., "rota": ...}
""" """
# Get worker names from the first row, starting from column 5 (index 4) # Get worker names from the first row, starting from column 5 (index 4)
worker_names = [str(x).strip() for x in calender_df.columns[5:]] worker_names = [str(x).strip() for x in calender_df.columns[5:]]
@@ -49,15 +62,11 @@ def extract_leave_and_rota_from_calender(calender_df):
work_requests = [] work_requests = []
# Process leave data from row 3 onwards (index 2 and up) # Process leave data from row 3 onwards (index 2 and up)
for idx, row in calender_df.iloc[3:].iterrows(): for idx, row in calender_df.iloc[3:].iterrows():
date_str = str(row.iloc[1]).strip().split(" ")[0] # Get the date string from the second column (index 1) date = _parse_sheet_date(row.iloc[1])
if not date_str or date_str.lower() == "nan": if date is None:
continue continue
try:
date = datetime.datetime.strptime(date_str, "%Y-%m-%d").date()
except Exception:
continue # skip rows with invalid date
if date < datetime.date(2026, 2, 16): if date < datetime.date.fromisoformat(ROTA_START_DATE):
continue # skip rows before rota start date continue # skip rows before rota start date
for i, worker in enumerate(worker_names): for i, worker in enumerate(worker_names):
if not worker or worker.lower() == "nan": if not worker or worker.lower() == "nan":
@@ -96,29 +105,30 @@ def extract_shift_assignments_from_rota(rota_df):
assignments = [] assignments = []
skipped_rows = 0 skipped_rows = 0
rota_start = datetime.date.fromisoformat(ROTA_START_DATE)
for idx, row in rota_df.iloc[2:].iterrows(): for idx, row in rota_df.iloc[2:].iterrows():
week_start_str = str(row.iloc[0]).strip().split(" ")[0] # Get the week start date from the first column (index 0) week_start = _parse_sheet_date(row.iloc[0])
try: if week_start is None:
week_start = datetime.datetime.strptime(week_start_str, "%Y-%m-%d").date() print(f"Invalid date format in row {idx}: {row.iloc[0]}")
except Exception:
print(f"Invalid date format in row {idx}: {week_start_str}")
continue # skip rows with invalid date continue # skip rows with invalid date
if week_start < datetime.date(2026, 2, 16): if week_start < rota_start:
skipped_rows += 1 skipped_rows += 1
continue # skip rows before rota start date continue # skip rows before rota start date
week = ((week_start - rota_start).days // 7) + 1
for i, day in enumerate(days): for i, day in enumerate(days):
for j, shift in enumerate(shifts): for j, shift in enumerate(shifts):
if i < 5: if i < 5:
# Mon-Fri: two columns per day (twilight, oncall) # Mon-Fri: two columns per day (twilight, oncall)
col_idx = 1 + i * 2 + j # 1 for Mon twilight, 2 for Mon oncall, etc. col_idx = 2 + i * 2 + j # 2 for Mon twilight, 3 for Mon oncall, etc.
weekend = False weekend = False
else: else:
# Sat/Sun: only one column for "weekend" shift (twilight), skip "night" # Sat/Sun: only one column for "weekend" shift (twilight), skip "night"
weekend = True weekend = True
if j == 0: if j == 0:
col_idx = 11 + (i - 5) # 11 for Sat, 12 for Sun col_idx = 12 + (i - 5) # 12 for Sat, 13 for Sun
else: else:
continue # No night shift column for Sat/Sun continue # No night shift column for Sat/Sun
if col_idx >= len(row): if col_idx >= len(row):
@@ -131,7 +141,7 @@ def extract_shift_assignments_from_rota(rota_df):
shift = "weekend" shift = "weekend"
assignments.append({ assignments.append({
"week_start": week_start, "week_start": week_start,
"week": idx-1-skipped_rows, # 1-indexed week number "week": week,
"day": day, "day": day,
"date": week_start + datetime.timedelta(days=i), "date": week_start + datetime.timedelta(days=i),
"shift": shift, "shift": shift,
@@ -149,7 +159,7 @@ def extract_shift_assignments_from_rota(rota_df):
def load_workers(): def load_workers():
# Path to the ODS file # Path to the ODS file
ods_path = "cons/CONSULTANT TWILIGHT & ONCALL ROTA.ods" ods_path = ROTA_REQUESTS
# Read all sheets # Read all sheets
xls = pd.read_excel(ods_path, engine="odf", sheet_name=None) xls = pd.read_excel(ods_path, engine="odf", sheet_name=None)
@@ -160,13 +170,19 @@ def load_workers():
rota_df = xls["Rota"] rota_df = xls["Rota"]
rota_b_path = "cons/Rota B Feb -July 2026.xlsx" rota_b_path = ROTA_B_PATH
rota_b_df = pd.read_excel(rota_b_path, engine="openpyxl", sheet_name="Sheet1", header=None) rota_b_df = pd.read_excel(rota_b_path, engine="openpyxl", sheet_name="Sheet1", header=None)
# Extract rota_b assignments: date in column A, worker in column B # Extract rota_b assignments: date in column A, worker in column B
rota_b_assignments = defaultdict(list) rota_b_assignments = defaultdict(list)
for idx, row in rota_b_df.iterrows(): for idx, row in rota_b_df.iterrows():
logger.debug(f"{idx}: {row}")
date_val = row.iloc[0] date_val = row.iloc[0]
worker_val = row.iloc[1].upper() if len(row) > 1 else None try:
worker_val = row.iloc[1].upper() if len(row) > 1 else None
except AttributeError:
# end of file
break
if pd.isna(date_val) or pd.isna(worker_val): if pd.isna(date_val) or pd.isna(worker_val):
continue continue
try: try:
@@ -236,19 +252,15 @@ def load_workers():
rota_data = {} rota_data = {}
for idx, row in days_df.iloc[7:].iterrows(): for idx, row in days_df.iloc[7:].iterrows():
# Extract initials from the name column (assumed to be in brackets at the end) # Extract initials from the name column (assumed to be in brackets at the end)
text = str(row.iloc[1]).strip() name = str(row.iloc[1]).strip()
if "(" in text and ")" in text: initial = str(row.iloc[2]).strip()
initial = text.split("(")[-1].split(")")[0].strip()
else:
raise ValueError(f"Invalid format in row {idx}: {text}")
name = text.split("(")[0].strip()
if not name or name.lower() == "nan": if not name or name.lower() == "nan":
continue continue
availability = {} availability = {}
requests = {} requests = {}
rota_data[name] = f"rota {str(row.iloc[2]).strip().lower()}" rota_data[name] = f"rota {str(row.iloc[3]).strip().lower()}"
for i, day in enumerate(weekday_cols): for i, day in enumerate(weekday_cols):
val = row.iloc[3 + i] val = row.iloc[4 + i]
# You can adjust the logic below depending on your marking scheme (e.g. "Y", "Yes", "1", etc.) # You can adjust the logic below depending on your marking scheme (e.g. "Y", "Yes", "1", etc.)
available = not ("no" in str(val).strip().lower() or str(val).strip().lower() == "yes*") available = not ("no" in str(val).strip().lower() or str(val).strip().lower() == "yes*")
availability[day] = available availability[day] = available
@@ -361,7 +373,8 @@ def load_workers():
worker = worker_day["name"] worker = worker_day["name"]
initial = worker_day["initial"] initial = worker_day["initial"]
start_date = "2026-02-16" # Default start date, can be adjusted later start_date = ROTA_START_DATE
#if initial == "ND": #if initial == "ND":
# start_date = "2025-09-29" # Non-working day worker starts later # start_date = "2025-09-29" # Non-working day worker starts later
@@ -404,8 +417,6 @@ def load_workers():
previous_shifts=priors_map.get(initial, {}), previous_shifts=priors_map.get(initial, {}),
) )
if initial == "LB":
w.start_date = datetime.date(2026, 6, 15)
#if initial == "L1": #if initial == "L1":
# w.oop = [OutOfProgramme( # w.oop = [OutOfProgramme(
@@ -447,19 +458,13 @@ def load_workers():
#w.add_force_assign_with("twilight", "oncall") #w.add_force_assign_with("twilight", "oncall")
w.add_force_assign_with("weekend", "oncall") w.add_force_assign_with("weekend", "oncall")
if w.name in ("DS",): #if w.name in ("DS",):
w.add_hard_day_dependency({"Fri", "Sat", "Sun"}, ignore_dates=[datetime.date(2026, 5, 2)]) # w.add_hard_day_dependency({"Fri", "Sat", "Sun"}, ignore_dates=[datetime.date(2026, 5, 2)])
else: #else:
w.add_hard_day_dependency({"Fri", "Sun"}) # w.add_hard_day_dependency({"Fri", "Sun"})
w.add_hard_day_exclusion({"Sat", "Sun"}) # w.add_hard_day_exclusion({"Sat", "Sun"})
#if w.name == "TB":
# w.add_hard_day_exclusion({"Sat", "Sun"}, start_date="2025-11-17")
#w.set_max_shifts
#w.set_max_shifts_per_week("weekend", 1)
workers.append(w) workers.append(w)
@@ -472,9 +477,9 @@ def main(
suspend: bool = False, suspend: bool = False,
solve: bool = True, solve: bool = True,
time_to_run: int = 60 * 60, time_to_run: int = 60 * 60,
ratio: float = 0.001, ratio: float = 0.1,
start_date: datetime.datetime = "2026-02-16", start_date: datetime.datetime = ROTA_START_DATE,
weeks: int = 24, weeks: int = 22,
bom: int = 1, bom: int = 1,
): ):
rota_start_date = start_date.date() rota_start_date = start_date.date()
@@ -485,7 +490,7 @@ def main(
weeks_to_rota=weeks, weeks_to_rota=weeks,
balance_offset_modifier=bom, balance_offset_modifier=bom,
use_previous_shifts=True, use_previous_shifts=True,
name="cons rota feb 2026 run 5", name="cons rota 2026 july",
allow_force_assignment_with_leave_conflict=True, allow_force_assignment_with_leave_conflict=True,
constraint_options=RotaConstraintOptions( constraint_options=RotaConstraintOptions(
balance_weekends=True, balance_weekends=True,
@@ -512,9 +517,9 @@ def main(
balance_offset=2, balance_offset=2,
assign_as_block=False, assign_as_block=False,
constraints=[ constraints=[
PreShiftConstraint(days=2, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), exclude_dates=["2026-04-06", "2026-05-04", "2026-05-25"]), #PreShiftConstraint(days=2, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), exclude_dates=["2026-04-06", "2026-05-04", "2026-05-25"]),
PreShiftConstraint(days=1, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), allow_self=False), #PreShiftConstraint(days=1, start_date="2025-11-17", exclude_days=("Fri", "Sat", "Sun"), allow_self=False),
MaxShiftsPerWeekConstraint(max_shifts=1, days=days[:5]), #MaxShiftsPerWeekConstraint(max_shifts=1, days=days[:5]),
], ],
), ),
SingleShift( SingleShift(
@@ -525,7 +530,7 @@ def main(
days=days[5:], days=days[5:],
balance_offset=1, balance_offset=1,
constraints=[ constraints=[
PostShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"]), #PostShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"]),
], ],
display_char="a", display_char="a",
#force_assign_with=["oncall"] #force_assign_with=["oncall"]
@@ -533,7 +538,7 @@ def main(
SingleShift( SingleShift(
sites=("rota b", "rota c"), sites=("rota b", "rota c"),
name="weekend b", name="weekend b",
end_date="2026-07-20", #end_date="2026-07-20",
workers_required=1, workers_required=1,
length=12.5, length=12.5,
days=days[5:], days=days[5:],
@@ -552,9 +557,9 @@ def main(
days=days[:5], days=days[:5],
balance_offset=1, balance_offset=1,
constraints=[ constraints=[
PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"], exclude_days=("Sat", "Sun")), #PreShiftConstraint(days=1, start_date="2025-11-17", ignore_shifts=["oncall"], exclude_days=("Sat", "Sun")),
MaxShiftsPerWeekConstraint(max_shifts=1), #MaxShiftsPerWeekConstraint(max_shifts=1),
MaxShiftsPerWeekBlockConstraint(week_block=3,max_shifts=1) #MaxShiftsPerWeekBlockConstraint(week_block=3,max_shifts=1)
], ],
), ),
) )