right-tree/backend/right_tree/api/management/commands/_spreadsheet_helpers.py

from openpyxl import load_workbook


def get_pk_mapping(object, mapping_key="name"):
    """ Returns a dictionary mapping a django model primary key to another given field.
    """
    pk_mapping = {}
    for instance in object.objects.all():
        pk_mapping[getattr(instance, mapping_key)] = instance.pk

    return pk_mapping


def get_col_mappings(sheet, start_col, row_index):
    """ Returns a dictionary that maps a spreadsheet cell value to a corresponding column index.
    """
    col_mappings = {}
    for row in sheet.iter_rows(min_col=start_col, min_row=row_index, max_row=row_index, values_only=True):
        for i, col_name in enumerate(row):
            col_mappings[col_name] = i

    return col_mappings


def get_pk_list_from_str(values_str, pk_mapping, fixes={}):
    """ Given a list of comma separated values from the spreadsheet. Returns a list of primary keys that
        correspond to the relevant values with any given mapping fixes applied.
    """
    pk_list = []
    for value in values_str.split(','):
        processed_value = value.lstrip().rstrip().replace(
            '_', ' ').replace('-', ' ').replace('’', '\'')

        # Applies any mapping adjustments between spreadsheet data and the database values
        if fixes and processed_value in fixes:
            processed_value = fixes[processed_value]

        # Adds the pk value for the value in the databse
        if processed_value in pk_mapping:
            pk_list.append(pk_mapping[processed_value])

    return pk_list


def get_spreadsheet(data_path, spreadsheet_filename):
    """ Returns a spreadsheet from a resources directory given the data path and
        spreadsheet filename.
    """
    spreadsheet_path = data_path / 'resources' / spreadsheet_filename
    workbook = load_workbook(filename=spreadsheet_path)
    return workbook.active