Wrapper on top of Python's csv module for reading/writing CSV sheets as dictionaries. Includes ability to lookup rows by a specific column, thereby allowing the data to be easily matched to external databases, such as QuickSchools or PipelineDeals. - csv_tools.py

csv_tools.py

#!/Library/Frameworks/Python.framework/Versions/2.7/bin/python

import csv


class CSV_IO(object):
    """class for reading/writing CSV objects
    can work standalone or as the backbone for CSVMatch"""

    def __init__(self, filepath):
        self.filepath = filepath
        self.values = []
        self.rows = []

        self.read()

    def read(self, val_funct=lambda val: val):
        """use val_funct to operate on all the values before as they are read in"""
        with open(self.filepath, 'rU') as f:
            raw_csv = csv.DictReader(f)
            for row in raw_csv:
                row = {key: val_funct(val) for key, val in row.iteritems()}
                self.rows.append(row)
                self.values += row.values()
            return

    def save(self, filepath=None):
        if not filepath:
            filepath = self.filepath
        with open(filepath, 'w') as f:
            writer = csv.DictWriter(f, self.rows[0].keys())
            writer.writeheader()
            for row in self.rows:
                writer.writerow(row)
            return


class CSVMatch(CSV_IO):

    def row_for_value(self, key, value):
        """returns a list of matching rows
        key = the column name on the CSV
        value = the value to match in that column
        """
        if not value or not (value in self.values): return

        match = None
        for row in self.rows:
            if row[key] == value:
                if match: raise MultipleMatchError()
                match = row
        return match

    def row_for_object(self, match_function, object):
        """like row_for_value, but allows for a more complicated match.
        match_function takes three parameters (vals, row, object) and return true/false
        """
        for row in self.rows:
            if match_function(row, object):
                return row


class MultipleMatchError(RuntimeError):
    pass