columns.py

"""This file contains code related to "Think Bayes",
by Allen B. Downey, available from greenteapress.com

Copyright 2014 Allen B. Downey
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html
"""

from __future__ import print_function, division

import csv


def read_csv(filename, constructor):
    """Reads a CSV file, returns the header line and a list of objects.

    filename: string filename
    """
    fp = open(filename)
    reader = csv.reader(fp)

    header = reader.next()
    names = [s.lower() for s in header]

    objs = [make_object(t, names, constructor) for t in reader]
    fp.close()

    return objs


def write_csv(filename, header, data):
    """Writes a CSV file

    filename: string filename
    header: list of strings
    data: list of rows
    """
    fp = open(filename, 'w')
    writer = csv.writer(fp)
    writer.writerow(header)

    for t in data:
        writer.writerow(t)
    fp.close()


def print_cols(cols):
    """Prints the index and first two elements for each column.

    cols: list of columns
    """
    for i, col in enumerate(cols):
        print(i, col[0], col[1])


def make_col_dict(cols, names):
    """Selects columns from a dataset and returns a map from name to column.

    cols: list of columns
    names: list of names
    """
    col_dict = {}
    for name, col in zip(names, cols):
        col_dict[name] = col
    return col_dict


def make_object(row, names, constructor):
    """Turns a row of values into an object.

    row: row of values
    names: list of attribute names
    constructor: function that makes the objects

    Returns: new object
    """
    obj = constructor()
    for name, val in zip(names, row):
        func = constructor.convert.get(name, int)
        try:
            val = func(val)
        except:
            pass
        setattr(obj, name, val)
    obj.clean()
    return obj