def test_addfieldusingcontext(): table1 = (('foo', 'bar'), ('A', 1), ('B', 4), ('C', 5), ('D', 9)) expect = (('foo', 'bar', 'baz', 'quux'), ('A', 1, None, 3), ('B', 4, 3, 1), ('C', 5, 1, 4), ('D', 9, 4, None)) def upstream(prv, cur, nxt): if prv is None: return None else: return cur.bar - prv.bar def downstream(prv, cur, nxt): if nxt is None: return None else: return nxt.bar - cur.bar table2 = addfieldusingcontext(table1, 'baz', upstream) table3 = addfieldusingcontext(table2, 'quux', downstream) ieq(expect, table3) ieq(expect, table3)
def test_addfieldusingcontext_stateful(): table1 = (('foo', 'bar'), ('A', 1), ('B', 4), ('C', 5), ('D', 9)) expect = (('foo', 'bar', 'baz', 'quux'), ('A', 1, 1, 5), ('B', 4, 5, 10), ('C', 5, 10, 19), ('D', 9, 19, 19)) def upstream(prv, cur, nxt): if prv is None: return cur.bar else: return cur.bar + prv.baz def downstream(prv, cur, nxt): if nxt is None: return prv.quux elif prv is None: return nxt.bar + cur.bar else: return nxt.bar + prv.quux table2 = addfieldusingcontext(table1, 'baz', upstream) table3 = addfieldusingcontext(table2, 'quux', downstream) ieq(expect, table3) ieq(expect, table3)