├── .gitignore ├── README.md ├── data.py ├── excel.py ├── generate_assignments.py ├── models.py ├── requirements.txt ├── solver.py ├── wheniwork.py └── wiw_upload.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.ipynb 2 | *.pyc 3 | .vscode 4 | *.json 5 | *.pickle 6 | *.csv 7 | *.xlsx 8 | *.pickle 9 | .DS_Store 10 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # shift-optimizer 2 | 3 | A shift optimizer, intended for use with small-to-medium sized groups. 4 | 5 | Built upon the fantastic [CP-SAT Solver](https://developers.google.com/optimization/cp/cp_solver) by Google. 6 | 7 | ## What it does 8 | Takes a list of work shifts, along with a list of personal preferences, and tries to find an optimal assignment for everyone. 9 | 10 | ## CLI 11 | Type 12 | ```shell 13 | $ python main.py -h 14 | ``` 15 | for help on how to use the program. 16 | -------------------------------------------------------------------------------- /data.py: -------------------------------------------------------------------------------- 1 | """Handling of raw data from files""" 2 | import json 3 | from solver import ShiftSolver 4 | from typing import List, Dict, Tuple 5 | from requests.auth import HTTPBasicAuth 6 | from datetime import datetime 7 | from models import Schedule, User, Shift, ShiftPreference 8 | import pytz 9 | 10 | def filter_unique_ordered(l): 11 | """Filter a list so that the items are unique. 12 | When an item appears more than once, the first occurrence will be retained. 13 | """ 14 | occured = set() 15 | filtered = [] 16 | for elem in l: 17 | if elem not in occured: 18 | occured.add(elem) 19 | filtered.append(elem) 20 | return filtered 21 | 22 | timestamp = int 23 | def datetime_string(ts: timestamp) -> str: 24 | return datetime.fromtimestamp(int(ts)).isoformat() 25 | 26 | def get_shifts(rshifts: List[Dict], timezone: str) -> List[Shift]: 27 | """Creates the necessary shift dict format 28 | Arguments: 29 | rshifts: rshifts 30 | timezone: timezone name 31 | Returns: 32 | list of Shifts 33 | """ 34 | shifts = list() 35 | for shift in rshifts: 36 | begin = datetime.fromtimestamp(int(float(shift['begin']))).astimezone(pytz.timezone(timezone)) 37 | end = datetime.fromtimestamp(int(float(shift['end']))).astimezone(pytz.timezone(timezone)) 38 | shifts.append( 39 | Shift( 40 | id=int(shift['id']), 41 | begin=begin, 42 | end=end, 43 | capacity=shift['capacity'], 44 | position=shift['position'] 45 | ) 46 | ) 47 | return shifts 48 | 49 | def get_users(rusers: List[Dict]) -> List[User]: 50 | """Get the schedule requirements for this person 51 | Arguments: 52 | users: [ 53 | { 54 | email 55 | hours_adjusted 56 | hours_max 57 | preferences 58 | positions 59 | } 60 | ] 61 | min_ratio: float, ratio of min hours to max hours 62 | Returns: 63 | list of users 64 | """ # WARNING highly custom code for Wolt Hungary 65 | users = list() 66 | for user in rusers: 67 | min_hours = user['hours_adjusted']**0.89 if user['hours_max'] >= 35 else 0.6 * user['hours_adjusted'] 68 | users.append( 69 | User( 70 | id=user['email'], 71 | min_hours=min_hours, 72 | max_hours=user['hours_adjusted'], 73 | only_long=(user['hours_max'] >= 35), # Fulltimer or not 74 | min_long=1, 75 | positions=user['positions'] 76 | ) 77 | ) 78 | return users 79 | 80 | def get_preferences(users: List[User], shifts: List[Shift], rusers: List[dict]) -> List[ShiftPreference]: 81 | # Index by id 82 | user = {u.id:u for u in users} 83 | shift = {s.id:s for s in shifts} 84 | preferences = [] 85 | for ruser in rusers: 86 | for rshiftid, priority in ruser['preferences'].items(): 87 | preferences.append(ShiftPreference( 88 | user=user[ruser['email']], 89 | shift=shift[int(rshiftid)], 90 | priority=priority 91 | )) 92 | return preferences 93 | 94 | def load_data(data: dict) -> Schedule: 95 | """ 96 | Args: 97 | data: { 98 | 'shifts': [ 99 | { 100 | 'id': int 101 | 'begin': timestamp 102 | 'end': timestamp 103 | 'capacity': int 104 | 'position': wiw_id 105 | } 106 | ] 107 | 'timezone': tzname 108 | 'users' [ 109 | { 110 | 'email': str 111 | 'hours_adjusted': float 112 | 113 | 'hours_max': float 114 | 'wiw_id': wiw_id 115 | 'preferences': { 116 | shiftid: priority 117 | } 118 | } 119 | ] 120 | } 121 | Returns: 122 | Schedule with the associated data""" 123 | rshifts = data['shifts'] 124 | rtimezone = data['timezone'] 125 | rusers = data['users'] 126 | shifts = get_shifts(rshifts, rtimezone) 127 | users = get_users(rusers) 128 | preferences = get_preferences(users, shifts, rusers) 129 | return Schedule(users,shifts,preferences) 130 | 131 | def json_compatible_solve(values: dict, data: dict) -> "dict[list,list]": 132 | """Create a solution object 133 | Args: 134 | values: [(day,shift,person)] = True | False 135 | data: { 136 | 'shifts': [ 137 | { 138 | 'id': int 139 | 'begin': timestamp 140 | 'end': timestamp 141 | 'capacity': int 142 | 'position': wiw_id 143 | } 144 | ] 145 | 'timezone': tzname 146 | 'users' [ 147 | { 148 | 'email': str 149 | 'hours_adjusted': float 150 | 'positions': [ 151 | position_id 152 | ] 153 | 'hours_max': float 154 | 'wiw_id': wiw_id 155 | 'preferences': { 156 | shiftid: priority 157 | } 158 | } 159 | ] 160 | } 161 | Returns: 162 | shifts: [ 163 | { 164 | user_id: wiw_user_id, 165 | position_id: wiw_pos_id, 166 | start_time: datetimestring 167 | end_time: datetimestring 168 | } 169 | ] 170 | """ 171 | user_wiw_for_email = dict() 172 | for user in data['users']: 173 | user_wiw_for_email[user['email']] = user['wiw_id'] 174 | shift_for_id = dict() 175 | capacity = dict() # capacity for shift id 176 | for s in data['shifts']: 177 | shift_for_id[s['id']] = { 178 | # Capacity is determined by number of shifts in solution 179 | 'start_time': datetime_string(s['begin']), 180 | 'end_time': datetime_string(s['end']), 181 | 'position_id': s['position'] 182 | } 183 | capacity[s['id']] = s['capacity'] 184 | # Add assigned shifts 185 | shifts = [] 186 | filled = {k: 0 for k in capacity.keys()} # n of capacities for shift id 187 | for (shift_id, user_email), assigned in values.items(): 188 | if assigned: 189 | shift = shift_for_id[shift_id].copy() 190 | shift['user_id'] = user_wiw_for_email[user_email] 191 | shifts.append(shift) 192 | filled[shift_id] += 1 193 | # Add openshifts 194 | for shift_id in capacity.keys(): 195 | # Add as many shifts to openshifts as many capacities are unfilled 196 | for _ in range(capacity[shift_id] - filled[shift_id]): 197 | shift = shift_for_id[shift_id].copy() 198 | shift['user_id'] = 0 199 | shift['is_shared'] = True 200 | shifts.append(shift) 201 | return shifts 202 | 203 | def solve_from_json_compatible(jsondict): 204 | """Create a solver values dict from 205 | a JSON-converted dict 206 | Args: 207 | assigned[day][shift][person] = True | False 208 | Returns: 209 | values: [(day,shift,person)] = True | False 210 | """ 211 | values = dict() 212 | for d, v1 in jsondict.items(): 213 | for s, v2 in v1.items(): 214 | for p, assigned_status in v2.items(): 215 | values[d,int(s),p] = assigned_status 216 | return values 217 | 218 | def empty_assignments(shifts, preferences): 219 | """Generate empty assignments 220 | Args: 221 | shifts: dict of sdata[day_id, shift_id] = { 222 | 'capacity': 2, 223 | 'begin': 525, 224 | 'end': 960 225 | } 226 | preferences: dict of pref[day_id,shift_id,person_id] = pref_score 227 | """ 228 | people = sorted(list(set([p for d,s,p in preferences.keys()]))) 229 | 230 | assignments = dict() 231 | 232 | for d, s in shifts.keys(): 233 | for p in people: 234 | assignments[d,s,p] = False 235 | 236 | return assignments 237 | 238 | def override_preqs(filename, preqs): 239 | """Use the .json file at the given path to override the personal reqs dict 240 | Args: 241 | filename: the JSON file to override with 242 | preqs: the preqs[person_id] = { 243 | 'min': n1, 244 | 'max': n2, 245 | 'min_long_shifts': n3, 246 | 'only_long_shifts': bool1 247 | } dict to override 248 | """ 249 | with open(filename, 'r', encoding='utf8') as jsonfile: 250 | override = json.load(jsonfile) 251 | 252 | for person in override: 253 | preqs[person] = override[person] 254 | 255 | def write_report(filename: str, solutions: List[Tuple[str,ShiftSolver]]): 256 | """Generate and write to file a report about the several solutions""" 257 | txt = '' 258 | for sol_file, sol in solutions: 259 | txt += f"""----- {sol_file} ----- 260 | Capacities filled: {sol.FilledCapacities}/{sol.NCapacities} ({round(sol.FilledCapacities/sol.NCapacities*100,2)}%) 261 | Hours filled: {sol.FilledHours}/{sol.Hours} ({round(sol.FilledHours/sol.Hours*100,2)}%) 262 | Prefscore: {sol.PrefScore} 263 | ------- 264 | """ 265 | with open(filename, 'w', encoding='utf8') as txtfile: 266 | txtfile.write(txt) -------------------------------------------------------------------------------- /excel.py: -------------------------------------------------------------------------------- 1 | import xlsxwriter 2 | from xlsxwriter.utility import xl_rowcol_to_cell as celln 3 | import data 4 | from colour import Color 5 | from typing import List, Tuple 6 | from solver import ShiftSolver 7 | 8 | def get_days(shifts): 9 | """Get the list of day names, in order, 10 | from a shift tuple list. 11 | """ 12 | unique_days = [] 13 | for d, s in shifts.keys(): 14 | del s 15 | if d not in unique_days: 16 | unique_days.append(d) 17 | return unique_days 18 | 19 | def get_prefcolor(n, max_n): 20 | """Takes a pref score, and a worst pref score, 21 | and returns a color to format how good that pref score is. 22 | Args: 23 | n: the pref score 24 | max_n: the worst pref score for that shift 25 | Returns: 26 | color: long hex color 27 | """ 28 | assert n <= max_n 29 | good = Color('#bef7c5') 30 | bad = Color('#ffd9d9') 31 | return list(good.range_to(bad, max_n + 1))[n].hex_l 32 | 33 | def get_people(preferences): 34 | return sorted(list(set([p for d,s,p in preferences.keys()]))) 35 | 36 | def write_to_file(filename, shifts, preferences, assignments, personal_reqs): 37 | """Write the solver result to an Excel workbook 38 | Args: 39 | shifts: dict of sdata[day_id, shift_id] = { 40 | 'capacity': 2, 41 | 'begin': 525, 42 | 'end': 960 43 | } 44 | prefs: dict of preferences[day_id,shift_id,person_id] = pref_score 45 | assignments: dict of assigned[day_id,shift_id,person_id] = True | False 46 | personal_reqs: dict of preqs[person_id] = { 47 | 'min': n1, 48 | 'max': n2, 49 | 'min_long_shifts': n3, 50 | 'only_long_shifts': bool1 51 | } 52 | """ 53 | workbook = xlsxwriter.Workbook(filename) 54 | shifts_ws = workbook.add_worksheet(name="shifts") 55 | time_f = workbook.add_format({'num_format': 'hh:mm'}) 56 | 57 | # Shifts 58 | n_shifts = len(shifts) 59 | ## Headers 60 | for idx, txt in enumerate(["Day", "ShiftID", "Capacity", "Begin", "End", "strID", "length in hours"]): 61 | shifts_ws.write(0, idx, txt) 62 | 63 | for rowidx, ((d,s), sdata) in enumerate(shifts.items(), start=1): 64 | shifts_ws.write(rowidx, 0, d) # Day 65 | shifts_ws.write(rowidx, 1, s) # ShiftId 66 | shifts_ws.write(rowidx, 2, sdata['capacity']) # Capacity 67 | shifts_ws.write(rowidx, 3, sdata['begin']/(24*60), time_f) # Begin time 68 | shifts_ws.write(rowidx, 4, sdata['end']/(24*60), time_f) # End time 69 | shifts_ws.write(rowidx, 5, str(d)+str(s)) # strID - should be unique 70 | shifts_ws.write_formula(rowidx, 6, f'=({celln(rowidx, 4, col_abs=True)}-{celln(rowidx, 3, col_abs=True)})*24') # shift length 71 | 72 | # Preferences 73 | people = get_people(preferences) 74 | n_people = len(people) 75 | percentage_format = workbook.add_format({'num_format': '0.00%'}) 76 | 77 | # Add default none values to preferences dict 78 | for d,s in shifts.keys(): 79 | for p in people: 80 | if (d,s,p) not in preferences.keys(): 81 | preferences[d,s,p] = None 82 | 83 | # Write to the sheet 84 | pref_ws = workbook.add_worksheet(name="preferences") 85 | pref_ws.protect() 86 | ## Headers 87 | for idx, txt in enumerate(["strID"] + people): 88 | pref_ws.write(0, idx, txt) 89 | for rowidx, (d,s) in enumerate(shifts.keys(), start=1): 90 | pref_ws.write(rowidx, 0, str(d)+str(s)) # strID 91 | for colidx, p in enumerate(people, start=1): 92 | pref_ws.write(rowidx, colidx, preferences[d,s,p]) 93 | 94 | ## % of shifts registered to last row 95 | pref_ws.write(n_shifts+1,0, "Registered%") 96 | for pidx in range(1, n_people + 1): 97 | pref_ws.write_formula(n_shifts + 1, pidx, 98 | ''.join(f'''=IF( 99 | NOT( 100 | VLOOKUP({celln(0, pidx)}, 101 | personal_reqs!{celln(1,0)}:{celln(n_people,4)}, 102 | 5) 103 | ), 104 | COUNT({celln(1,pidx)}:{celln(n_shifts,pidx)})/{n_shifts}, 105 | COUNT({celln(1,pidx)}:{celln(n_shifts,pidx)})/COUNTIF(assignments!{celln(1,n_people+3)}:{celln(n_shifts,n_people+3)},"Long") 106 | ) 107 | '''.split()), cell_format=percentage_format) 108 | 109 | 110 | # Personal requirements 111 | pers_ws = workbook.add_worksheet(name="personal_reqs") 112 | ## Headers 113 | for idx, txt in enumerate(["person","Min. hours", "Max. hours","Min. long shits","Only long shifts"]): 114 | pers_ws.write(0, idx, txt) 115 | for rowidx, p in enumerate(people, start=1): 116 | pers_ws.write(rowidx, 0, p) 117 | pers_ws.write_number(rowidx, 1, personal_reqs[p]['min']) 118 | pers_ws.write_number(rowidx, 2, personal_reqs[p]['max']) 119 | pers_ws.write_number(rowidx, 3, personal_reqs[p]['min_long_shifts']) 120 | pers_ws.write_boolean(rowidx, 4, personal_reqs[p]['only_long_shifts']) 121 | 122 | # Assignments 123 | 124 | # Formats 125 | no_pref = workbook.add_format({'font_color':'#dedede'}) 126 | 127 | assign_ws = workbook.add_worksheet(name="assignments") 128 | # Write to the sheet 129 | ## Headers 130 | for idx, txt in enumerate(["strID"] + people + ["n assigned", "capacity", "Is long"]): 131 | assign_ws.write(0, idx, txt) 132 | 133 | for rowidx, (d,s) in enumerate(shifts.keys(), start=1): 134 | assign_ws.write(rowidx, 0, str(d)+str(s)) # strID 135 | for colidx, p in enumerate(people, start=1): 136 | if preferences[d,s,p] is None: 137 | assign_ws.write_boolean(rowidx, colidx, assignments[d,s,p], no_pref) 138 | else: 139 | pref_color = get_prefcolor(preferences[d,s,p], max([preferences[d,s,p] for p in people if preferences[d,s,p] is not None])) 140 | pref_format = workbook.add_format({'bg_color':pref_color}) 141 | assign_ws.write_boolean(rowidx, colidx, assignments[d,s,p], pref_format) 142 | # Add formula t calculate number of empty shifts 143 | assign_ws.write(n_shifts+1,n_people+1, 'Empty shifts') 144 | assign_ws.write_formula(n_shifts+1,n_people+2, f'=COUNTIF({celln(1,n_people+1)}:{celln(n_shifts,n_people+1)},0)') 145 | # Add formula to calculate number of empty places in shifts 146 | assign_ws.write(n_shifts+2,n_people+1, 'Empty places on shifts') 147 | assign_ws.write_formula(n_shifts+2,n_people+2, f'=SUM({celln(1,n_people+2)}:{celln(n_shifts,n_people+2)})-SUM({celln(1,n_people+1)}:{celln(n_shifts,n_people+1)})') 148 | # Add total pref score 149 | assign_ws.write(n_shifts+3,n_people+1, 'Pref score') 150 | assign_ws.write(n_shifts+3,n_people+2,f'=SUM({celln(n_shifts+7,1)}:{celln(n_shifts+7,n_people)})') 151 | 152 | # Add shift capacity condition indicator 153 | ## Add a formula to the end of each row, 154 | ## to calculate the number of people assigned. 155 | ## Add conditional formatting to this cell, 156 | ## so that it's: 157 | ## green, when the shift is full 158 | ok_format = workbook.add_format({'bg_color':'#a6ff9e'}) 159 | ## orange, when the shift is below capacity 160 | medium_format = workbook.add_format({'bg_color': '#ffc670'}) 161 | ## red, when the shift is over capacity 162 | bad_format = workbook.add_format({'bg_color': '#ff7a70'}) 163 | for rowidx, (d,s) in enumerate(shifts.keys(), start=1): 164 | # Current shift state 165 | cap_format = workbook.add_format({'left':1}) 166 | assign_ws.write_formula( 167 | rowidx, n_people+1, 168 | f'=COUNTIF({celln(rowidx, 1)}:{celln(rowidx,n_people)}, TRUE)', 169 | cap_format) 170 | # Capacity 171 | assign_ws.write_formula( 172 | rowidx, n_people+2, 173 | f'=shifts!C{rowidx+1}') 174 | assign_ws.write_formula( 175 | rowidx, n_people+3, 176 | f'=IF(INDEX(shifts!{celln(1,0)}:{celln(n_shifts,6)},MATCH(assignments!{celln(rowidx,0)},assignments!{celln(1,0)}:{celln(n_shifts,0)},0),7)>5,"Long", "")' 177 | ) 178 | 179 | #region Conditional formatting the n_people on shift 180 | capacity_col_row = celln(rowidx, n_people+2) 181 | # Full 182 | assign_ws.conditional_format(celln(rowidx, n_people+1), 183 | { 184 | 'type':'cell', 185 | 'criteria':'==', 186 | 'value': capacity_col_row, 187 | 'format': ok_format 188 | }) 189 | # Below capacity 190 | assign_ws.conditional_format(celln(rowidx, n_people+1), 191 | { 192 | 'type':'cell', 193 | 'criteria':'<', 194 | 'value': capacity_col_row, 195 | 'format': medium_format 196 | }) 197 | # Full 198 | assign_ws.conditional_format(celln(rowidx, n_people+1), 199 | { 200 | 'type':'cell', 201 | 'criteria':'>', 202 | 'value': capacity_col_row, 203 | 'format': bad_format 204 | }) 205 | #endregion 206 | 207 | # Add personal requirement indicator 208 | ## For each person, under their last shift, indicate their: 209 | ## Minimum hours (from the personal_reqs sheet) 210 | ## The actual number of hours they are taking 211 | ## Maximum hours (from the personal_reqs sheet) 212 | ## Number of long shifts taken 213 | ## Minimum number of long shifts (from the personal_reqs sheet) 214 | ## Whether they only take long shifts (from the personal_reqs sheet) 215 | 216 | # Horrifying excel formula #1 217 | def preq_formula(person_col, var_col): 218 | return (f'=INDEX(personal_reqs!{celln(1,0, row_abs=True, col_abs=True)}:{celln(n_people,4, row_abs=True, col_abs=True)},'+ 219 | f'MATCH({celln(0,person_col, row_abs=True)},personal_reqs!{celln(1,0, row_abs=True, col_abs=True)}:{celln(n_people, 0, row_abs=True, col_abs=True)},0)' 220 | +f',{var_col})') 221 | 222 | def workhours_formula(person_col): 223 | return (f'=SUMIF({celln(1,person_col, row_abs=True)}:{celln(n_shifts, person_col, row_abs=True)},TRUE,shifts!{celln(1,6, row_abs=True, col_abs=True)}:{celln(n_shifts,6, row_abs=True, col_abs=True)})') 224 | 225 | def pref_score_formula(person_col): 226 | return f'=SUMIF(assignments!{celln(1,person_col, row_abs=True)}:{celln(n_shifts,person_col, row_abs=True)},TRUE,preferences!{celln(1,person_col, row_abs=True)}:{celln(n_shifts,person_col, row_abs=True)})' 227 | 228 | # TODO minor formatting here 229 | 230 | # Row headers 231 | for row_idx, txt in enumerate(["Min. hours", 232 | "Actual hours", 233 | "Max. hours", 234 | "Long shifts taken", 235 | "Min. long shifts", 236 | "Long shifts only", 237 | "Pref score" 238 | ]): 239 | r0 = n_shifts+1 240 | assign_ws.write(r0+row_idx, 0, txt) 241 | 242 | for col_idx, p in enumerate(people, start=1): 243 | r0 = n_shifts+1 # Start after the last shift 244 | min_hours_formula = preq_formula(col_idx, 2) 245 | assign_ws.write_formula(r0, col_idx, min_hours_formula) 246 | assign_ws.write_formula(r0+1, col_idx, workhours_formula(col_idx)) 247 | max_hours_formula = preq_formula(col_idx, 3) 248 | assign_ws.write_formula(r0+2, col_idx, max_hours_formula) 249 | assign_ws.write_formula(r0+3, col_idx, f'=COUNTIFS({celln(1,col_idx)}:{celln(n_shifts,col_idx)},TRUE,{celln(1, n_people+3)}:{celln(n_shifts, n_people+3)},"Long")') 250 | min_long_shifts_formula = preq_formula(col_idx, 4) 251 | assign_ws.write_formula(r0+4, col_idx, min_long_shifts_formula) 252 | long_shifts_only_formula = preq_formula(col_idx, 5) 253 | assign_ws.write_formula(r0+5, col_idx, long_shifts_only_formula) 254 | assign_ws.write_formula(r0+6, col_idx, pref_score_formula(col_idx)) 255 | 256 | # Add conditional to show if hours is withing range 257 | assign_ws.conditional_format(celln(r0+1,col_idx), 258 | { 259 | 'type':'cell', 260 | 'criteria':'between', 261 | 'minimum': celln(r0,col_idx), 262 | 'maximum': celln(r0+2,col_idx), 263 | 'format': ok_format 264 | }) 265 | assign_ws.conditional_format(celln(r0+1,col_idx), 266 | { 267 | 'type':'cell', 268 | 'criteria':'not between', 269 | 'minimum': celln(r0,col_idx), 270 | 'maximum': celln(r0+2,col_idx), 271 | 'format': bad_format 272 | }) 273 | 274 | # Add conditional to highlight actual applications 275 | applied_shift_format = workbook.add_format({ 276 | 'bold': True, 277 | 'border': 1 278 | }) 279 | 280 | assign_ws.conditional_format(0,0, n_shifts+1, n_people+1, 281 | { 282 | 'type': 'cell', 283 | 'criteria': '==', 284 | 'value': True, 285 | 'format': applied_shift_format 286 | } 287 | ) 288 | 289 | # Master view 290 | master = workbook.add_worksheet(name='sensei-view') 291 | for col_idx, txt in enumerate(["strID", 292 | "Begin", 293 | "End", 294 | "Person", 295 | "Works" 296 | ]): 297 | master.write(0, col_idx, txt) # TODO hide works columns 298 | 299 | for scount, ((d,s), sdata) in enumerate(shifts.items()): 300 | for pcount, person in enumerate(people): 301 | r_index = scount*n_people+pcount + 1 302 | master.write(r_index, 0, str(d)+str(s)) 303 | master.write(r_index, 1, sdata['begin']/(24*60), time_f) 304 | master.write(r_index, 2, sdata['end']/(24*60), time_f) 305 | master.write(r_index, 3, person) 306 | assignments_range = f'assignments!{celln(1,1)}:{celln(n_shifts, n_people)}' 307 | shiftname_addr = f'{celln(r_index, 0)}' 308 | shiftids_range = f'assignments!{celln(1,0)}:{celln(n_shifts,0)}' 309 | pname_addr = f'{celln(r_index, 3)}' 310 | names_range = f'assignments!{celln(0,1)}:{celln(0, n_people)}' 311 | master.write_formula(r_index, 4, f'=INDEX({assignments_range},MATCH({shiftname_addr},{shiftids_range},0),MATCH({pname_addr},{names_range},0))') 312 | 313 | master.autofilter(0,4,0,4) 314 | master.filter_column(4, "Works == TRUE") 315 | 316 | # Worker view 317 | worker = workbook.add_worksheet(name='padawan-view') 318 | for col_idx, txt in enumerate([ 319 | "Person", 320 | "strID", 321 | "Begin", 322 | "End", 323 | "Works" 324 | ]): 325 | worker.write(0, col_idx, txt) # TODO hide works columns 326 | for pcount, person in enumerate(people): 327 | for scount, ((d,s), sdata) in enumerate(shifts.items()): 328 | r_index = pcount*n_shifts+scount + 1 329 | worker.write(r_index, 0, person) 330 | worker.write(r_index, 1, str(d)+str(s)) 331 | worker.write(r_index, 2, sdata['begin']/(24*60), time_f) 332 | worker.write(r_index, 3, sdata['end']/(24*60), time_f) 333 | assignments_range = f'assignments!{celln(1,1)}:{celln(n_shifts, n_people)}' 334 | shiftname_addr = f'{celln(r_index, 1)}' 335 | shiftids_range = f'assignments!{celln(1,0)}:{celln(n_shifts,0)}' 336 | pname_addr = f'{celln(r_index, 0)}' 337 | names_range = f'assignments!{celln(0,1)}:{celln(0, n_people)}' 338 | worker.write_formula(r_index, 4, f'=INDEX({assignments_range},MATCH({shiftname_addr},{shiftids_range},0),MATCH({pname_addr},{names_range},0))') 339 | 340 | worker.autofilter(0,0,0,4) # For the your name 341 | worker.filter_column(4, "Works == TRUE") 342 | 343 | 344 | workbook.close() 345 | 346 | def write_summary(filename: str, rows: Tuple[str, ShiftSolver]): 347 | """Creates an excel worksheet to a new file showing the properties of the solves, 348 | with links to them. 349 | Args: 350 | filename: to create the workbook at 351 | rows: list of (xlsxfilepath, solver) tuples 352 | """ 353 | workbook = xlsxwriter.Workbook(filename) 354 | 355 | hour_format = workbook.add_format({'num_format': '#,##0.00'}) 356 | 357 | ws = workbook.add_worksheet('index') 358 | for cidx, txt in enumerate( 359 | ['Prefscore', 360 | 'Average prefscore/person', 361 | 'Mode of prefscores', 362 | 'Empty Shifts', 363 | 'Unfilled Hours', 364 | 'Unfilled Hours (%)', 365 | 'Unfilled Capacities', 366 | 'Unfilled Capacities (%)', 367 | 'Marginal prefence cost of capacity (f\')', 368 | 'Link to Solve']): 369 | ws.write(0, cidx, txt) 370 | 371 | percentage_format = workbook.add_format({'num_format': '0.00%'}) 372 | dec_format = workbook.add_format({'num_format': '0.00'}) 373 | is_first_row = True 374 | for rowidx, (solutionpath, solver) in enumerate(rows, start=1): 375 | for colidx, (val, format_) in enumerate( 376 | [(solver.PrefScore(),None), 377 | (f'={(c_pref := celln(rowidx,0))}/{(c_n_all_people := celln(4,11))}', dec_format), 378 | (",".join(list(map(str,solver.PrefModes()))), None), 379 | (solver.EmptyShifts(), None), 380 | (solver.UnfilledHours(), hour_format), 381 | (f'={(c_hours := celln(rowidx, 4))}/{(c_n_all_hours := celln(3,11))}', percentage_format), 382 | (solver.UnfilledCapacities(), None), 383 | (f'={(c_caps := celln(rowidx, 6))}/{(c_n_all_caps := celln(2,11))}', percentage_format), 384 | (f'=IF(ISNUMBER({(c_caps_prev := celln(rowidx-1, 6))}),IF(AND((({c_caps_prev}-{c_caps})<>0),(({(c_pref_prev := celln(rowidx-1,0))}-{c_pref})<>0)),(-({c_pref}-{c_pref_prev})/({c_caps}-{c_caps_prev})),0),"")', None) 385 | ] 386 | ): 387 | ws.write(rowidx, colidx, val, format_) 388 | ws.write_url(rowidx, colidx+1, solutionpath) 389 | 390 | marginal_cost_chart = workbook.add_chart({'type':'line'}) 391 | marginal_cost_chart.add_series({'values': f'=index!{celln(0, 8)}:{celln(len(rows), 8)}'}) 392 | marginal_cost_chart.set_title({'name': 'Marginal cost for each solve'}) 393 | marginal_cost_chart.set_legend({'none': True}) 394 | ws.insert_chart('M3', marginal_cost_chart) 395 | 396 | try: 397 | # Use the first solution, these values should be equal everywhere 398 | solver = rows[0][1] 399 | ws.write(1, 10, "Number of shifts") 400 | ws.write_number(1,11, solver.NShifts()) 401 | ws.write(2, 10, "Number of capacities") 402 | ws.write_number(c_n_all_caps, solver.NCapacities()) 403 | ws.write(3, 10, "Number of hours") 404 | ws.write_number(c_n_all_hours, solver.Hours()) 405 | ws.write(4, 10, "Number of people") 406 | ws.write_number(c_n_all_people, solver.NPeople()) 407 | except IndexError: 408 | # No solutions 409 | pass 410 | 411 | workbook.close() 412 | -------------------------------------------------------------------------------- /generate_assignments.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import data 3 | import json 4 | from solver import ShiftSolver 5 | import excel 6 | from pathlib import Path 7 | from copy import deepcopy 8 | 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument('file', 11 | help='Path to the .json file with the schedule data.') 12 | 13 | parser.add_argument('--no-solve', dest='nosolve', 14 | help="Don't solve, just create the overview excel.", action='store_true') 15 | 16 | parser.add_argument('-t', '--timeout', 17 | help='The maximum time in seconds that the solver can take to find an optimal solution.', default=None, type=int) 18 | 19 | parser.add_argument('-c', '--capacities', 20 | help='The percentage of capacities to fill as a minimum', default=96.0, type=float) 21 | 22 | parser.add_argument('-f', '--force-availabilities', dest='force_available', 23 | help='Extend shift availability for every position for each user.', action='store_true') 24 | args = parser.parse_args() 25 | 26 | with open(args.file, 'r') as f: 27 | jsondata = json.load(f) 28 | schedule = data.load_data(jsondata) 29 | 30 | if args.force_available: # Optionally extend solution space 31 | schedule.add_forced_availabilities() 32 | 33 | solver = ShiftSolver(schedule) 34 | 35 | sum_capacities = 0 36 | # Calculate the number of capacities total 37 | for shift in schedule.shifts: 38 | sum_capacities += shift.capacity 39 | 40 | starting_capacity = int(sum_capacities*(args.capacities / 100)) 41 | 42 | subfolderpath = '.' 43 | 44 | Path(subfolderpath+'/sols').mkdir(parents=True, exist_ok=True) 45 | rows = [] # {'pref':s1, 'unfilled':s2, 'empty':s3, 'filename':s4} 46 | 47 | if not args.nosolve: 48 | for n in range(starting_capacity, sum_capacities+1): 49 | if solver.Solve( 50 | timeout=args.timeout, 51 | min_capacities_filled=n): 52 | print(f'Prefscore: {solver.ObjectiveValue()} Unfilled capacities: {solver.UnfilledCapacities} in {round(solver.WallTime(),2)} seconds', end='') 53 | if solver.StatusName() != 'OPTIMAL': 54 | print(' !SUBOPTIMAL SOLVE! Try to run with more time', end='') 55 | print() 56 | filename = f'{n}.json' 57 | # Write to excel and add index for the root later 58 | rows.append((filename, deepcopy(solver))) 59 | 60 | with open(f'{subfolderpath}/sols/{n}.json', 'w', encoding='utf8') as jsonfile: 61 | json.dump(data.json_compatible_solve(solver.Values, jsondata), jsonfile, indent=4, ensure_ascii=False) 62 | else: # No more solutions to be found 63 | break 64 | if len(rows) > 0: 65 | data.write_report(f'{subfolderpath}/sols/solindex.txt', rows) 66 | -------------------------------------------------------------------------------- /models.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime, date, timedelta, tzinfo, time 2 | from typing import List, Dict, Tuple, Set, Any, NewType 3 | UserId = NewType('UserId', Any) 4 | ShiftId = NewType('ShiftId', int) 5 | class Shift: 6 | """Shift timeframe and capacity 7 | Sorts based on begin time 8 | """ 9 | def __init__(self, id: ShiftId, begin: datetime, end: datetime, capacity: int, position: int): 10 | self.id = id 11 | assert begin < end 12 | self.begin = begin 13 | self.end = end 14 | self.capacity = capacity 15 | self.position = position 16 | @property 17 | def length(self) -> timedelta: 18 | return self.end - self.begin 19 | @property 20 | def is_long(self) -> bool: 21 | return self.length > timedelta(hours=6) 22 | @property 23 | def starts_early(self) -> bool: 24 | return self.begin.time() < time(9, 00) 25 | @property 26 | def ends_late(self) -> bool: 27 | return self.end.time() > time(22,00) or self.begin.date() < self.end.date() 28 | def __eq__(self, other: "Shift"): 29 | return ((self.begin, self.end, self.capacity) == (other.begin, other.end, other.capacity)) 30 | def __ne__(self, other: "Shift"): 31 | return not self.__eq__(other) 32 | def __lt__(self, other: "Shift"): 33 | return self.begin < other.begin 34 | def __le__(self, other: "Shift"): 35 | return self.begin <= other.begin 36 | def __gt__(self, other: "Shift"): 37 | return self.begin > other.begin 38 | def __ge__(self, other: "Shift"): 39 | return self.begin >= other.begin 40 | def __and__(self, other: "Shift"): 41 | """Detects if two shifts overlap""" 42 | return ((self.begin <= other.begin <= self.end) or 43 | (self.begin <= other.end <= self.end) or 44 | (other.begin <= self.begin <= other.end) or 45 | (other.begin <= self.end <= other.end)) 46 | 47 | def __repr__(self): 48 | return f'{self.begin}-{self.end} Capacity {self.capacity}' 49 | class User: 50 | """User requirements""" 51 | def __init__(self, id: UserId, positions: List[int], min_hours: float, max_hours: float, only_long: bool, min_long: int): 52 | self.id = id 53 | self.positions = positions 54 | assert min_hours < max_hours 55 | self.min_hours = min_hours 56 | self.max_hours = max_hours 57 | self.only_long = only_long 58 | self.min_long = min_long 59 | self._availabilities = None 60 | def can_take(self, shift: Shift) -> bool: 61 | return shift.position in self.positions and (not self.only_long or shift.is_long) 62 | def set_availabilities_for(self, schedule): 63 | """Extract a set of time intervals 64 | in which the user is available for work, 65 | indicated by the fact that they've taken at least a single shift in that span, 66 | and set it to a member variable, accessible through a property. 67 | """ 68 | self._availabilities = set() 69 | for preference in schedule.preferences: 70 | if preference.user == self: 71 | self._availabilities.add((preference.shift.begin, preference.shift.end)) 72 | @property 73 | def availabilities(self) -> Set[Tuple[datetime,datetime]]: 74 | if self._availabilities is None: 75 | raise ValueError("Availabilities accessed before having been set") 76 | return self._availabilities 77 | def is_available_at(self, shift: Shift) -> bool: 78 | """Checks whether the user can take a shift, 79 | based on whether they have the necessary position, 80 | and if they're available at that time. 81 | """ 82 | for begin, end in self.availabilities: 83 | if begin <= shift.begin <= end and begin <= shift.end <= end: 84 | return True # User is available for this shift 85 | return False 86 | 87 | class ShiftPreference: 88 | """Stores a User-Shift relation with a preference score""" 89 | def __init__(self, user: User, shift: Shift, priority: int): 90 | self.user = user 91 | self.shift = shift 92 | self.priority = priority 93 | def __eq__(self, other: "ShiftPreference") -> bool: 94 | return ( 95 | self.user == other.user and 96 | self.shift == other.shift 97 | ) # Ignore priority when checking equality 98 | class Schedule: 99 | """Schedule information""" 100 | def __init__(self, users: List[User], shifts: List[Shift], preferences: List[ShiftPreference]): 101 | self.users = users 102 | self.shifts = shifts 103 | self.preferences = preferences 104 | self._shifts_for_day = None 105 | self.user = {u.id:u for u in users} # index id 106 | self.shift = {s.id:s for s in shifts} # index id 107 | self._preference = None 108 | @property 109 | def shifts_for_day(self) -> Dict[date, List[Shift]]: 110 | """Collects shifts for a given day for each day, 111 | in order of start time 112 | Returns: 113 | dict[date] = [s1, s2, ...] 114 | """ 115 | if self._shifts_for_day is not None: 116 | return self._shifts_for_day # cached result 117 | # Collect days in order 118 | days = [] 119 | for s in self.shifts: 120 | day = s.begin.date() 121 | if day not in days: 122 | days.append(day) 123 | days.sort() 124 | self._shifts_for_day = {day:[] for day in days} 125 | # Collect shifts for day 126 | for shift in self.shifts: 127 | day = shift.begin.date() 128 | self._shifts_for_day[day].append(shift) 129 | # Sort 130 | for lst in self._shifts_for_day.values(): 131 | lst.sort() 132 | return self._shifts_for_day 133 | @property 134 | def preference(self) -> Dict[Tuple[ShiftId, UserId], int]: 135 | """Collect priority number for each shift and user""" 136 | if self._preference is not None: 137 | return self._preference # cached result 138 | self._preference = dict() 139 | for pref in self.preferences: 140 | self._preference[pref.shift.id, pref.user.id] = pref.priority 141 | return self._preference 142 | def add_forced_availabilities(self): 143 | """Create valid ShiftPreferences for cases where 144 | the user is available at the time of a shift 145 | (meaning they signed up for another shift in that timeframe), 146 | and they can take the shift based on their positions, 147 | but they didn't originally sign up for that shift. 148 | """ 149 | self._preference = None # Force recalculate preference cache 150 | for user in self.users: 151 | user.set_availabilities_for(schedule=self) # Calculate availability intervals 152 | for shift in self.shifts: 153 | if user.can_take(shift) and user.is_available_at(shift): 154 | sp = ShiftPreference( 155 | user=user, 156 | shift=shift, 157 | priority=100 # very low priority 158 | ) 159 | if sp not in self.preferences: 160 | # Avoid overriding existing preferences 161 | self.preferences.append(sp) -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | colour==0.1.5 2 | ortools==8.0.8283 3 | XlsxWriter==1.3.7 4 | -------------------------------------------------------------------------------- /solver.py: -------------------------------------------------------------------------------- 1 | from ortools.sat.python import cp_model 2 | from itertools import combinations, permutations 3 | from models import Schedule, Shift, ShiftId, User, ShiftPreference 4 | from typing import List, Dict, Any, NoReturn, Tuple, Set, Iterable, Optional 5 | from datetime import timedelta 6 | 7 | class ShiftModel(cp_model.CpModel): 8 | """Shift solver 9 | 10 | This class takes a list of shifts, 11 | and a list of preferenced shift requests. 12 | Then provides an optimal solution (if one exists) 13 | for the given parameters. 14 | """ 15 | def __init__(self, schedule: Schedule): 16 | """Args: 17 | schedule: the Schedule to solver for 18 | """ 19 | super().__init__() 20 | self.schedule = schedule 21 | self.variables = { # Create solver variables 22 | (p.shift.id, p.user.id):self.NewBoolVar(f'{p.user.id} works {p.shift.id}') 23 | for p in self.schedule.preferences 24 | } 25 | # Add must-have-constraints 26 | self.AddNoConflict() 27 | 28 | def AddMaxDailyShifts(self, n: int = 1): 29 | """Make sure that employees only get assigned to 30 | maximum of n shifts on any given day. 31 | Args: 32 | n: the max number of shifts a person can work on any day. 33 | """ 34 | for u in self.schedule.users: 35 | for shifts_for_day in self.schedule.shifts_for_day.values(): 36 | self.AddLinearConstraint(sum([self.variables[s.id,u.id] for s in shifts_for_day if (s.id,u.id) in self.variables]), 0, n) 37 | 38 | def AddShiftCapacity(self): 39 | """Make sure that no more people are assigned to a shift than its capacity""" 40 | for s in self.schedule.shifts: 41 | self.AddLinearConstraint( 42 | sum([self.variables[s.id,u.id] for u in self.schedule.users 43 | if (s.id,u.id) in self.variables]), 0, s.capacity) 44 | 45 | def AddMinimumCapacityFilledNumber(self, n: int): 46 | """Make sure that at least n out of the sum(capacities) is filled. 47 | """ 48 | self.Add(sum([assigned_val for assigned_val in self.variables.values()]) >= n) 49 | 50 | def AddMinMaxWorkTime(self): 51 | """Make sure that everyone works within their schedule time range. 52 | """ 53 | for u in self.schedule.users: 54 | worktime = 0 55 | for s in self.schedule.shifts: 56 | if (s.id,u.id) in self.variables: 57 | worktime += self.variables[s.id, u.id] * s.length.seconds 58 | self.AddLinearConstraint(worktime, int(u.min_hours*60*60), int(u.max_hours*60*60)) 59 | 60 | def AddLongShifts(self): 61 | """Make sure that everyone works at least n long shifts. 62 | Args: 63 | n: the number of shifts one needs to work 64 | """ 65 | for u in self.schedule.users: 66 | if u.only_long: 67 | self.Add( 68 | sum([self.variables[s.id,u.id] for s in self.schedule.shifts if not s.is_long]) == 0 69 | ) 70 | elif u.min_long > 0: 71 | self.Add( 72 | sum([self.variables[s.id,u.id] for s in self.schedule.shifts if s.is_long]) > u.min_long 73 | ) 74 | 75 | def AddLongShiftBreak(self): 76 | """Make sure that if you work a long shift, you're not gonna work 77 | another shift on the same day. 78 | """ 79 | for u in self.schedule.users: 80 | for s in self.schedule.shifts: 81 | if s.is_long and (s.id,u.id) in self.variables: 82 | # Technically: for each long shift, if p works on that long shift, 83 | # Make sure that for that day, 84 | # The number of shifts worked for that person is exactly one. 85 | self.Add(sum([( 86 | self.variables[other_s.id,u.id]) 87 | for other_s in self.schedule.shifts_for_day[s.begin.date()] 88 | if (other_s.id,u.id) in self.variables] 89 | ) == 1).OnlyEnforceIf(self.variables[s.id, u.id]) 90 | 91 | def AddNoConflict(self): 92 | """Make sure that no one has two shifts on a day that overlap. 93 | """ 94 | conflicting_pairs = set() # assuming every day has the same shifts 95 | for shift1, shift2 in combinations(self.schedule.shifts, r=2): 96 | if shift1 & shift2: # bitwise and -> overlaps 97 | conflicting_pairs.add((shift1.id,shift2.id)) 98 | 99 | for u in self.schedule.users: 100 | for s1_id, s2_id in conflicting_pairs: 101 | # Both of them can't be assigned to the same person 102 | # => their sum is less than 2 103 | if (s1_id,u.id) in self.variables and (s2_id,u.id) in self.variables: 104 | self.Add(self.variables[s1_id,u.id] + self.variables[s2_id,u.id] < 2) 105 | 106 | def AddSleep(self): 107 | """Make sure that no one has a shift in the morning, 108 | if they had a shift last evening. 109 | """ 110 | conflicting_pairs = self.get_nosleep_shifts() 111 | 112 | for u in self.schedule.users: 113 | for s1, s2 in conflicting_pairs: 114 | if (s1,u.id) in self.variables and (s2,u.id) in self.variables: 115 | # Both of them can't be assigned to the same person 116 | # => their sum is less than 2 117 | self.Add(self.variables[s1,u.id] + self.variables[s2, u.id] < 2) 118 | 119 | def AddNonFulltimerMaxShifts(self, n: int): 120 | """Make sure that non-fulltimers (people who can take non-long shifts) 121 | Can take at most n shifts. 122 | """ 123 | for u in self.schedule.users: 124 | if not u.only_long: 125 | self.AddLinearConstraint( 126 | sum([self.variables[s.id,u.id] for s in self.schedule.shifts if (s.id,u.id) in self.variables]), 127 | 0, 128 | n) 129 | 130 | def MaximizeWelfare(self): 131 | """Maximize the welfare of the employees. 132 | This target will minimize the dissatisfaction of the employees 133 | with their assigned shift. 134 | """ 135 | self.Minimize( 136 | sum([works*self.schedule.preference[shift,user] for (shift, user), works in self.variables.items()]) 137 | ) 138 | 139 | # Helper methods 140 | def get_nosleep_shifts(self) -> Set[Tuple[Shift,Shift]]: 141 | """Collect pairs of shifts that conflict in the following way: 142 | The time between the end of one and the begin of the other is 143 | less than 11 hours for long shifts, and 9 hours for non-long. 144 | """ 145 | conflicting = set() 146 | for shift, other_shift in permutations(self.schedule.shifts, r=2): 147 | if shift.end < other_shift.begin: 148 | offtime_between = other_shift.begin - shift.end 149 | if shift.is_long: 150 | if timedelta(0) <= offtime_between <= timedelta(hours=11): 151 | conflicting.add((shift.id, other_shift.id)) 152 | else: 153 | if timedelta(0) <= offtime_between <= timedelta(hours=9): 154 | conflicting.add((shift.id, other_shift.id)) 155 | return conflicting 156 | 157 | class ShiftSolver(cp_model.CpSolver): 158 | def __init__(self, schedule: Schedule): 159 | """Args: 160 | schedule: the Schedule to solve for 161 | """ 162 | super().__init__() 163 | self.schedule=schedule 164 | self.__model = None 165 | 166 | def Solve(self, min_capacities_filled: int = 0, timeout: Optional[int]=None) -> bool: 167 | """ 168 | Args: 169 | min_workers: The minimum number of workers that have to be assigned to every shift 170 | hours: hours[person_id] = {'min': n1, 'max': n2} dict 171 | n_long_shifts: Number of long shifts for every worker 172 | pref_function: function that takes and returns an integer, used for weighting of the pref function 173 | timeout: number of seconds that the solver can take to find the optimal solution 174 | Returns: 175 | Boolean: whether the solver found a solution. 176 | """ 177 | 178 | self.__model = ShiftModel(self.schedule) 179 | self.__model.AddMinimumCapacityFilledNumber(n=min_capacities_filled) 180 | self.__model.MaximizeWelfare() 181 | self.__model.AddShiftCapacity() 182 | self.__model.AddLongShiftBreak() 183 | self.__model.AddSleep() 184 | self.__model.AddMinMaxWorkTime() 185 | self.__model.AddMaxDailyShifts(1) 186 | self.__model.AddNonFulltimerMaxShifts(5) 187 | if timeout is not None: 188 | self.parameters.max_time_in_seconds = timeout 189 | super().Solve(self.__model) 190 | if super().StatusName() in ('FEASIBLE', 'OPTIMAL'): 191 | return True 192 | return False 193 | 194 | def get_overview(self): 195 | return self.get_shift_workers() + self.get_employees_hours() 196 | 197 | def get_shift_workers(self): 198 | """Human-readable overview of the shifts 199 | Returns: 200 | Multiline string 201 | """ 202 | txt = '' 203 | for shift in self.__model.schedule.shifts: 204 | txt += f'{shift}' 205 | txt += ''.join( 206 | [f'\n\t{u.id} p={self.__model.schedule.preference[shift.id,u.id]}' for u in self.schedule.users 207 | if (shift.id,u.id) in self.__model.variables 208 | and self.Value(self.__model.variables[shift.id,u.id])]) 209 | txt += '\n' 210 | return txt 211 | 212 | def get_employees_hours(self): 213 | """Human-readable hours for each employee 214 | Returns: 215 | Multiline string 216 | """ 217 | txt = str() 218 | for u in self.__model.schedule.users: 219 | work_hours=0 220 | for shift in self.__model.schedule.shifts: 221 | if (shift.id,u.id) in self.__model.variables: 222 | work_hours += self.Value(self.__model.variables[shift.id,u.id]) * shift.length.seconds / (60*60) 223 | txt += f'{u.id} works {round(u.min_hours, 2)}<={round(work_hours, 2)}<={round(u.max_hours, 2)} hours.\n' 224 | return txt 225 | 226 | @property 227 | def Values(self) -> dict: 228 | """Returns a dictionary with the solver values. 229 | Returns: 230 | assigned[shift_id, person_id] = True | False 231 | """ 232 | assigned = dict() 233 | for s,p in self.__model.variables.keys(): 234 | assigned[s,p] = self.Value(self.__model.variables[s,p]) 235 | 236 | return assigned 237 | 238 | @property 239 | def PrefScore(self) -> float: 240 | return self.ObjectiveValue() 241 | 242 | @property 243 | def NShifts(self) -> int: 244 | return len(self.__model.schedule.shifts) 245 | 246 | @property 247 | def UnfilledCapacities(self) -> int: 248 | assigned = self.Values 249 | unfilled_capacities = 0 250 | for s in self.__model.schedule.shifts: 251 | unfilled_capacities += (s.capacity- sum([assigned[s.id,u.id] for u in self.__model.schedule.users if (s.id,u.id) in assigned])) 252 | return unfilled_capacities 253 | 254 | @property 255 | def FilledCapacities(self) -> int: 256 | return self.NCapacities - self.UnfilledCapacities 257 | @property 258 | def NCapacities(self) -> int: 259 | capacities = 0 260 | for shift in self.__model.schedule.shifts: 261 | capacities += shift.capacity 262 | return capacities 263 | @property 264 | def UnfilledHours(self) -> float: 265 | assigned = self.Values 266 | unfilled_hours = 0 267 | for shift in self.__model.schedule.shifts: 268 | unfilled_capacities_on_this_shift = (shift.capacity - sum([assigned[shift.id,u.id] for u in self.__model.schedule.users if (shift.id,u.id) in assigned])) 269 | length_of_shift_in_hours = shift.length.seconds / (60*60) 270 | unfilled_hours += unfilled_capacities_on_this_shift * length_of_shift_in_hours 271 | return unfilled_hours 272 | @property 273 | def FilledHours(self) -> float: 274 | return self.Hours - self.UnfilledHours 275 | @property 276 | def Hours(self) -> float: 277 | n_hours = 0 278 | for shift in self.__model.schedule.shifts: 279 | n_hours += shift.length.seconds / (60*60) 280 | return n_hours 281 | 282 | @property 283 | def NPeople(self) -> int: 284 | return len(self.__model.people) 285 | -------------------------------------------------------------------------------- /wheniwork.py: -------------------------------------------------------------------------------- 1 | """ 2 | Wrapper for the WhenIWork API 3 | 4 | The get methods return the exact JSON response, 5 | unless stated otherwise. 6 | """ 7 | 8 | from datetime import datetime 9 | import requests 10 | from typing import List 11 | import json 12 | 13 | class WhenIWorkError(Exception): pass 14 | 15 | class NoTokenError(WhenIWorkError): pass 16 | 17 | class NoLoginError(WhenIWorkError): pass 18 | 19 | class WhenIWork: 20 | def __init__(self, token, user_id): 21 | self.token = token 22 | if token is None: 23 | raise NoTokenError("User doesn't have token specified") 24 | self.user_id = user_id 25 | def __get(self, address, params={}) -> dict: 26 | """ 27 | Send a GET request to WhenIWork, 28 | and returns the result as a dictionary 29 | Args: 30 | params (dict): the parameters to pass along 31 | Returns: 32 | response (dict): the response dictionary 33 | """ 34 | headers= {"W-Token": self.token, "W-UserId": str(self.user_id)} 35 | response = requests.get( 36 | f'https://api.wheniwork.com/2/{address}', 37 | headers=headers, 38 | params=params 39 | ) 40 | response.raise_for_status() # throw if not 200 41 | return response.json() 42 | def __post(self, address, data={}): 43 | """ 44 | Send a POST request to WhenIWork. 45 | Args: 46 | data (dict): the data to send 47 | """ 48 | headers = {"W-Token": self.token, "W-UserId": str(self.user_id)} 49 | response = requests.post( 50 | f'https://api.wheniwork.com/2/{address}', 51 | headers=headers, 52 | data=data 53 | ) 54 | response.raise_for_status() 55 | return response.json() 56 | @classmethod 57 | def get_token(cls, api_key: str, email: str, password: str) -> dict: 58 | response = requests.post( 59 | 'https://api.login.wheniwork.com/login', 60 | headers={ 61 | "W-Key": api_key, 62 | 'content-type': 'application/json' 63 | }, 64 | data='{"email":"'+email+'","password":"'+password+'"}', 65 | ) 66 | response.raise_for_status() 67 | return response.json() 68 | def get_locations(self, only_unconfirmed=False) -> "list[dict]": 69 | """ 70 | Get the locations associated with this workplace 71 | Args: 72 | only_unconfirmed (bool): Include only unconfirmed schedules/locations 73 | Returns: 74 | resp (list): array of location objects 75 | """ 76 | params = {'only_unconfirmed': only_unconfirmed} 77 | locations = self.__get('locations', params=params)['locations'] 78 | return locations 79 | def create_location(self, params: dict): 80 | """Create Schedule(Location) 81 | https://apidocs.wheniwork.com/external/index.html#tag/Schedules-(Locations)/paths/~12~1locations/post 82 | Arguments: 83 | params: { 84 | account_id: int 85 | name: str 86 | address: str 87 | coordinates: [float, float] 88 | deleted_at: str 89 | ip_address: str 90 | is_default: bool 91 | is_deleted: bool 92 | latitude: float 93 | longitude: float 94 | max_hours: int 95 | place: { 96 | address: str 97 | business_name: str 98 | country: str, 99 | id: int, 100 | latitude: float, 101 | locality: str, 102 | longitude: float, 103 | place_id: str 104 | place_type: [], 105 | postal_code: [str,...] 106 | region: str, 107 | street_name": str, 108 | street_number": str, 109 | sub_locality": str, 110 | updated_at": str 111 | } 112 | place_confirmed: bool 113 | place_id: str 114 | radius: int 115 | sort: int 116 | created_at: str 117 | updated_at: str 118 | } 119 | """ 120 | return self.__post('locations', data=params) 121 | def get_users(self, show_pending=True, show_deleted=False, search=None) -> dict: 122 | """ 123 | Get users from the workplace 124 | Returns: 125 | resp (dict): array under key `users` 126 | """ 127 | params = { 128 | 'show_pending': show_pending, 129 | 'show_deleted': show_deleted, 130 | 'search': search 131 | } 132 | return self.__get('users', params=params) 133 | def get_positions(self, show_deleted=False) -> dict: 134 | """ 135 | Get positions from the workplace 136 | Returns: 137 | resp (dict): array under key `positions` 138 | """ 139 | params = { 'show_deleted': show_deleted } 140 | return self.__get('positions', params=params) 141 | def get_timeoff_types(self) -> dict: 142 | """ 143 | Get Time Off Types 144 | Returns: 145 | resp (dict): array under key `request-types` 146 | """ 147 | return self.__get('requesttypes') 148 | def __get_timeoff_requests_pagination( 149 | self, 150 | start: datetime, 151 | end: datetime, 152 | user_id: int = None, 153 | location_id: int =None, 154 | max_id: int = None, 155 | limit: int = 5, 156 | page: int = 0, 157 | since_id: int = None, 158 | sortby: str = None, 159 | include_deleted_users: bool = False, 160 | type: int = None 161 | ) -> dict: 162 | """ 163 | Get Time Off Requests 164 | Used for pagination 165 | Returns: 166 | resp (dict): array under key `requests` 167 | """ 168 | params = { 169 | 'start': start, 170 | 'end': end, 171 | 'user_id': user_id, 172 | 'location_id': location_id, 173 | 'max_id': max_id, 174 | 'limit': limit, 175 | 'page': page, 176 | 'since_id': since_id, 177 | 'sortby': sortby, 178 | 'include_deleted_users': include_deleted_users, 179 | 'type': type 180 | } 181 | return self.__get('requests', params=params) 182 | def get_timeoff_requests( 183 | self, 184 | start: datetime, 185 | end: datetime, 186 | user_id: int = None, 187 | location_id: int =None, 188 | max_id: int = None, 189 | since_id: int = None, 190 | sortby: str = None, 191 | include_deleted_users: bool = False, 192 | type: int = None 193 | ) -> "list[dict]": 194 | """ 195 | Get all Time Off request objects 196 | Returns: 197 | resp (list): array of requests 198 | """ 199 | timeoff = [] 200 | total = 201 # default total, will be reset 201 | page = 0 202 | while page * 200 < total: 203 | resp = self.__get_timeoff_requests_pagination( 204 | start, 205 | end, 206 | user_id=user_id, 207 | location_id=location_id, 208 | max_id=max_id, 209 | since_id=since_id, 210 | sortby=sortby, 211 | include_deleted_users=include_deleted_users, 212 | type=type, 213 | limit=200, 214 | page=page 215 | ) 216 | page += 1 217 | total = resp['total'] # Make sure the total is right 218 | timeoff += resp['requests'] # Add request objects 219 | return timeoff 220 | def get_availabilities(self, start:datetime=None, end:datetime=None, user_id=None, include_all=None) -> dict: 221 | """ 222 | Returns: 223 | resp (dict): array under the key `availabilityevents` 224 | """ 225 | params = { 226 | 'start': start, 227 | 'end': end, 228 | 'user_id': user_id, 229 | 'include_all': include_all 230 | } 231 | return self.__get('availabilityevents', params=params) 232 | def get_shifts(self, start:datetime, end:datetime, unpublished: bool=False, **kwargs) -> "List[dict]": 233 | kwargs.update( 234 | { 235 | 'start': start.isoformat(), 236 | 'end': end.isoformat(), 237 | 'unpublished': unpublished 238 | } 239 | ) 240 | return self.__get('shifts', params=kwargs) 241 | def create_shift(self, location, start:datetime or str, end:datetime or str, **kwargs): 242 | if type(start) is datetime: start = start.isoformat() 243 | if type(end) is datetime: end = end.isoformat() 244 | kwargs.update( 245 | { 246 | 'location_id': location, 247 | 'start_time': start, 248 | 'end_time': end 249 | } 250 | ) 251 | return self.__post('shifts', data=json.dumps(kwargs)) 252 | 253 | -------------------------------------------------------------------------------- /wiw_upload.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from wheniwork import WhenIWork, NoLoginError 3 | import data 4 | import json 5 | from requests import HTTPError 6 | from pathlib import Path 7 | import pickle 8 | import os.path 9 | 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('filename') 12 | parser.add_argument('--email') 13 | parser.add_argument('--password') 14 | parser.add_argument('--userid') 15 | parser.add_argument('--apikey') 16 | args = parser.parse_args() 17 | 18 | tokenfile_path = 'wiwtoken.pickle' 19 | 20 | wiwcreds = None 21 | # wiwtoken.pickle stores the WIW token. 22 | # If it exists, assume we've already got a token, and use that to make requests. 23 | if os.path.exists(tokenfile_path): 24 | with open(tokenfile_path, 'rb') as wiwtokenfile: 25 | wiwcreds = pickle.load(wiwtokenfile) 26 | else: # Authenticate manually using password 27 | if None in (args.email, args.password, args.userid, args.apikey): 28 | raise NoLoginError() # Arguments are not specified 29 | wiwcreds = WhenIWork.get_token(args.apikey, args.email, args.password) 30 | wiwcreds['user_id'] = args.userid 31 | with open(tokenfile_path, 'wb') as wiwtokenfile: 32 | pickle.dump(wiwcreds, wiwtokenfile) 33 | 34 | account_id = wiwcreds['person']['id'] 35 | 36 | # We've authenticated, time to make requests. 37 | wiw = WhenIWork(wiwcreds['token'], wiwcreds['user_id']) 38 | users = wiw.get_users() 39 | location_id = wiw.get_users()['locations'][0]['id'] # Assume there's just one 40 | 41 | # Load shifts to upload 42 | with open(args.filename, 'r') as shiftfile: 43 | shifts = json.load(shiftfile) 44 | 45 | uploaded = [] 46 | failed = [] 47 | 48 | # Upload shifts 49 | for shift in shifts: 50 | try: 51 | uploaded_shift = wiw.create_shift( 52 | location=location_id, 53 | start=shift['start_time'], 54 | end=shift['end_time'], 55 | user_id=shift['user_id'], 56 | position_id=shift['position_id'] 57 | ) # Create new shift 58 | uploaded.append((shift, uploaded_shift)) 59 | except HTTPError as e: 60 | print(str(e)) 61 | failed.append((shift, e)) 62 | 63 | # Stats 64 | if len(uploaded) > 0: 65 | print('Successfully uploaded shift ids:') 66 | for shift, upload in uploaded: 67 | print(upload['shift']['id']) 68 | 69 | if len(failed) > 0: 70 | print('Failed to upload:') 71 | for shift, error in failed: 72 | print(json.dumps(shift)+' Error: '+ str(error)) 73 | --------------------------------------------------------------------------------