aggregate_racking.py

# aggregate_racking.py
#
# Convert tabular racking data with one core section and location per row
# to a single row per location with all corresponding core sections.
#
# Input:
# CSV with no headers and data rows of form [Core ID, Location ID].
#
# Output:
# CSV with a single row for each unique Location ID followed by all
# Core IDs corresppnding to that location.

import csv, os, sys

# Read UTF-8 encoded CSV file with two columns, section and location.
# Return list of (section, location) tuples.
def read_core_location_csv(cl_file: str) -> 'list[tuple[str, str]]':
    core_loc_data = []
    with open(cl_file, encoding='utf-8-sig') as f:
        reader = csv.reader(f)
        core_loc_data = [(core, loc) for core, loc in reader]
    return core_loc_data

def aggregate_core_locations(cl_data: 'list[tuple[str, str]]') -> 'dict[str, list]':
    agg = {}
    for core, loc in cl_data:
        if loc not in agg:
            agg[loc] = [core]
        else:
            agg[loc].append(core)
    return agg

def write_aggregate_csv(agg_file, agg_dict):
    with open(agg_file, 'wt', encoding='utf-8-sig') as f:
        writer = csv.writer(f)
        for loc, cores in agg_dict.items() :
            writer.writerow([loc] + cores)

def convert_dirs(input_dir, output_dir):
    csv_files = [f for f in os.listdir(input_dir) if f.endswith('.csv')]
    for cf in csv_files:
        print(f'Aggregating racking data in {cf}...')
        core_loc_data = read_core_location_csv(os.path.join(input_dir, cf))
        agg_dict = aggregate_core_locations(core_loc_data)
        agg_file = f"{cf.split('.')[0]}_aggregated.csv"
        write_aggregate_csv(os.path.join(output_dir, agg_file), agg_dict)

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage:\npython aggregate_racking.py [dir containing input .csv files] [dir to write .csv output files]")
    else:
        convert_dirs(sys.argv[1], sys.argv[2])