source: workshop2010/trunk/src/grokworkshop/mergecsv.py @ 17503

"""Merge two CSV files and output the results.
"""
import csv

def merge_csv_files(csv_path1=None, csv_path2=None,
                    wanted_headers=['id', 'student_id', 'name']):
    """Merge a file 'input1.csv' and 'input2.csv'.
    """
    input1 = csv_path1
    input2 = csv_path2

    students_reader = csv.DictReader(open(input1, 'rb'))
    payments_reader = csv.DictReader(open(input2, 'rb'))

    result_writer = None

    students = dict()
    for student in students_reader:
        # Read all students and store them under their respective id
        students[student['id']] = student
    for payment in payments_reader:
        # Read all payments and lookup a matching student
        student_id = payment['student_id']

        if student_id not in students.keys():
            # Oops, what happened?
            print "STUDENT NOT FOUND: ", student_id
            continue

        student = students[student_id]
        result_line = payment
        result_line.update(student)
        if result_writer is None:
            # Create a CSV writer and ouput header line
            result_writer = csv.DictWriter(open('result.csv', 'wb'), 
                                           wanted_headers, 
                                           extrasaction='ignore')
            header_line = dict([(x, x) for x in wanted_headers])
            result_writer.writerow(header_line)
        result_writer.writerow(result_line)