zscore xml fixture generator script

gistfile1.py

import csv
import os.path
import sys

from lxml import etree
from lxml.builder import ElementMaker

E = ElementMaker()

def loadcsv(path):
    with open(path) as f:
        r = csv.reader(f)
        for i, row in enumerate(r):
            if i == 0:
                continue

            age_month = int(row[0])
            weight_thresholds = [float(k) for k in row[4:]]
            yield {'age': age_month, 'thresh': weight_thresholds}

def e_weights(thresholds):
    LIMIT = 4
    for i in range(LIMIT):
        z = str(i - 3)
        lo = str(thresholds[i - 1] if i > 0 else 0)
        hi = str(thresholds[i]) if i < LIMIT - 1 else 'none'
        yield E.weight(z, low=lo, high=hi)

def e_month(m):
    return E.month(*e_weights(m['thresh']), id=str(m['age']))

def e_gender(gender, table):
    return E.gender(*(e_month(m) for m in table), id=gender)

dir = '/home/drew/tmp/'
tables = {
    'm': os.path.join(dir, 'zscore_m.csv'),
    'f': os.path.join(dir, 'zscore_f.csv'),
}

data = dict((gender, list(loadcsv(path))) for gender, path in tables.iteritems())

root = E.zscore(*(e_gender(gender, table) for gender, table in data.iteritems()))

etree.ElementTree(element=root).write(sys.stdout, encoding='utf-8', pretty_print=True)