assign_peers.py

from collections import defaultdict

INITIAL_SLOTS = 3
INITIAL_WORKLOAD = 3

def _assign_peer_reviewers(df):
    reviewer_slots = defaultdict(list)
    reviewer_workload = defaultdict(list)

    df_copy = df.copy()
    if "available_slots" not in df_copy.columns.values:
        df_copy["available_slots"] = INITIAL_SLOTS
    if "available_workload" not in df_copy.columns.values:
        df_copy["available_workload"] = INITIAL_WORKLOAD

    for row in df_copy.itertuples():
        employee = row.employee
        reviewer = row.collaborator_employee
        available_slots = row.available_slots
        available_workload = row.available_workload

        assigned_reviewer_slots = len(reviewer_slots[employee])
        assigned_reviewer_workload = len(reviewer_workload[employee])

        if (assigned_reviewer_slots < available_slots) & (assigned_reviewer_workload < available_workload):
            reviewer_slots[employee].append(reviewer)
            reviewer_workload[employee].append(employee)
        else:
            continue
    
    reviewer_slots_df = pd.DataFrame.from_dict(reviewer_slots, orient="index").T.unstack().dropna().reset_index(level=0)
    reviewer_slots_df = reviewer_slots_df.rename(columns={"level_0": "employee", 0: "collaborator_employee"})

    return reviewer_slots_df

team_member_reviewers_df = _assign_peer_reviewers(ranked_team_df)