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

from openpyxl import load_workbook
from right_tree.api.models import Zone


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_zone_pk_mapping(zone_model):
    """ Maps the string instance (unique) of a zone to its corresponding primary key.
    """
    zone_pk_mapping = {}
    for instance in zone_model.objects.all():
        zone_pk_mapping[str(instance)] = instance.pk

    return zone_pk_mapping


def get_col_mappings(sheet, start_col, stop_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, max_col=stop_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_zone_col_mappings(sheet, start_col, stop_col, start_row, stop_row):
    """ Returns a dictionary that maps a spreadsheet zone string object to a column index.
    """
    zone_col_mappings = {}
    current_zone = current_variant = current_refined_variant = None
    for i, col in enumerate(sheet.iter_cols(min_col=start_col, max_col=stop_col, min_row=start_row, max_row=stop_row, values_only=True)):
        zone_name, zone_variant, zone_refined_variant = col

        if zone_name is not None:
            current_zone = zone_name
            current_variant = current_refined_variant = None

        if zone_variant is not None:
            current_variant = zone_variant
            current_refined_variant = None

        current_refined_variant = zone_refined_variant if zone_refined_variant is not None else current_refined_variant

        zone_obj = Zone(name=current_zone, variant=current_variant,
                        refined_variant=current_refined_variant)
        zone_col_mappings[str(zone_obj)] = i + start_col

    return zone_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