Think Stats 1.4 -- high on functions

first.py

from datetime import timedelta
from functools import partial, wraps
from itertools import ifilter, imap
from operator import attrgetter, not_
import survey


def caller(method_name, *args, **kwargs):
    def call_obj(obj):
        return getattr(obj, method_name)(*args, **kwargs)
    return call_obj


def counts(xs):
    c = {}
    for x in xs:
        c.setdefault(x, 0)
        c[x] += 1
    return c


def compose(*functions):
    if not functions:
        raise ValueError("Must specify functions to compose")
    def composed(*args, **kwargs):
        fs = list(functions)
        y = fs.pop()(*args, **kwargs)
        while fs:
            f = fs.pop()
            y = f(y)
        return y
    return composed


def wrap_result(wrapper):
    return lambda f: wraps(f)(compose(wrapper, f))


negate = wrap_result(not_)


odd = lambda x: x % 2 == 1
even = negate(odd)


def on_items(f, d):
    return compose(dict, f, caller('items'))(d)


def map_dict(f, d):
    return on_items(partial(imap, f), d)


def filter_dict(p, d):
    return on_items(partial(ifilter, p), d)


def map_keys(f, d):
    return map_dict(lambda (k, v): (f(k), v), d)


def dichotomy(p, xs):
    return ifilter(negate(p), xs), ifilter(p, xs)


OUTCOME_LABELS = {
    1: 'LIVE BIRTH',
    2: 'INDUCED ABORTION',
    3: 'STILLBIRTH',
    4: 'MISCARRIAGE',
    5: 'ECTOPIC PREGNANCY',
    6: 'CURRENT PREGNANCY',
    }


get_pregnancy_lengths = partial(imap, attrgetter('prglength'))

def average(xs):
    xs = list(xs)
    if not xs:
        return 'EMPTY'
    return sum(xs) / float(len(xs))


table = survey.Pregnancies()
table.ReadRecords()
print 'Number of pregnancies', len(table.records)
live_births = filter(lambda record: record.outcome == 1, table.records)
print 'Live births:', len(live_births)
rest, firstborns = dichotomy(lambda record: record.birthord == 1, live_births)
[rest, firstborns] = map(compose(average, get_pregnancy_lengths), [rest, firstborns])
print 'Average pregnancy length (firstborns):', firstborns
print 'Average pregnancy length (others):', rest
print 'Difference:', timedelta(weeks=firstborns) - timedelta(weeks=rest)