def test_todb_appenddb_cursor():
f = NamedTemporaryFile(delete=False)
conn = sqlite3.connect(f.name)
conn.execute("create table foobar (foo, bar)")
conn.commit()
# exercise function
table = (("foo", "bar"), ("a", 1), ("b", 2), ("c", 2))
cursor = conn.cursor()
todb(table, cursor, "foobar")
# check what it did
actual = conn.execute("select * from foobar")
expect = (("a", 1), ("b", 2), ("c", 2))
ieq(expect, actual)
# try appending
table2 = (("foo", "bar"), ("d", 7), ("e", 9), ("f", 1))
appenddb(table2, cursor, "foobar")
# check what it did
actual = conn.execute("select * from foobar")
expect = (("a", 1), ("b", 2), ("c", 2), ("d", 7), ("e", 9), ("f", 1))
ieq(expect, actual)
def test_fromdb_withargs():
# initial data
data = (('a', 1),
('b', 2),
('c', 2.0))
connection = sqlite3.connect(':memory:')
c = connection.cursor()
c.execute('create table foobar (foo, bar)')
for row in data:
c.execute('insert into foobar values (?, ?)', row)
connection.commit()
c.close()
# test the function
actual = fromdb(
connection,
'select * from foobar where bar > ? and bar < ?',
(1, 3)
)
expect = (('foo', 'bar'),
('b', 2),
('c', 2.0))
ieq(expect, actual)
ieq(expect, actual) # verify can iterate twice
def test_unflatten():
table1 = (('lines',),
('A',),
(1,),
(True,),
('C',),
(7,),
(False,),
('B',),
(2,),
(False,),
('C',),
(9,))
expect1 = (('f0', 'f1', 'f2'),
('A', 1, True),
('C', 7, False),
('B', 2, False),
('C', 9, None))
actual1 = unflatten(table1, 'lines', 3)
ieq(expect1, actual1)
ieq(expect1, actual1)
def test_tosqlite3_appendsqlite3_connection():
conn = sqlite3.connect(':memory:')
# exercise function
table = (('foo', 'bar'),
('a', 1),
('b', 2),
('c', 2))
tosqlite3(table, conn, 'foobar', create=True)
# check what it did
actual = conn.execute('select * from foobar')
expect = (('a', 1),
('b', 2),
('c', 2))
ieq(expect, actual)
# check appending
table2 = (('foo', 'bar'),
('d', 7),
('e', 9),
('f', 1))
appendsqlite3(table2, conn, 'foobar')
# check what it did
actual = conn.execute('select * from foobar')
expect = (('a', 1),
('b', 2),
('c', 2),
('d', 7),
('e', 9),
('f', 1))
ieq(expect, actual)
def test_fromcsv_gz():
"""Test the fromcsv function on a gzipped file."""
f = NamedTemporaryFile(delete=False)
f.close()
fn = f.name + '.gz'
os.rename(f.name, fn)
fz = gzip.open(fn, 'wb')
writer = csv.writer(fz, delimiter='\t')
table = (('foo', 'bar'),
('a', 1),
('b', 2),
('c', 2))
for row in table:
writer.writerow(row)
fz.close()
actual = fromcsv(fn, delimiter='\t')
expect = (('foo', 'bar'),
('a', '1'),
('b', '2'),
('c', '2'))
ieq(expect, actual)
ieq(expect, actual) # verify can iterate twice
def test_sort_4():
table = (("foo", "bar"), ("C", 2), ("A", 9), ("A", 6), ("F", 1), ("D", 10))
result = sort(table, "bar")
expectation = (("foo", "bar"), ("F", 1), ("C", 2), ("A", 6), ("A", 9), ("D", 10))
ieq(expectation, result)
def test_fromxml_2():
# initial data
f = NamedTemporaryFile(delete=False)
data = """<table>
<tr>
<td v='foo'/><td v='bar'/>
</tr>
<tr>
<td v='a'/><td v='1'/>
</tr>
<tr>
<td v='b'/><td v='2'/>
</tr>
<tr>
<td v='c'/><td v='2'/>
</tr>
</table>"""
f.write(data)
f.close()
print open(f.name).read()
actual = fromxml(f.name, 'tr', 'td', 'v')
print actual
expect = (('foo', 'bar'),
('a', '1'),
('b', '2'),
('c', '2'))
ieq(expect, actual)
ieq(expect, actual) # verify can iterate twice
def test_extendheader():
table1 = (("foo",), ("a", 1, True), ("b", 2, False))
table2 = extendheader(table1, ["bar", "baz"])
expect2 = (("foo", "bar", "baz"), ("a", 1, True), ("b", 2, False))
ieq(expect2, table2)
ieq(expect2, table2) # can iterate twice?
def test_pushheader():
table1 = (("a", 1), ("b", 2))
table2 = pushheader(table1, ["foo", "bar"])
expect2 = (("foo", "bar"), ("a", 1), ("b", 2))
ieq(expect2, table2)
ieq(expect2, table2) # can iterate twice?
def test_records_shortrows():
"""Test the records function on a table with short rows."""
table = (('foo', 'bar'), ('a', 1), ('b',))
actual = records(table)
expect = ({'foo': 'a', 'bar': 1}, {'foo': 'b', 'bar': None})
ieq(expect, actual)
def test_skip():
table1 = (("#aaa", "bbb", "ccc"), ("#mmm",), ("foo", "bar"), ("a", 1), ("b", 2))
table2 = skip(table1, 2)
expect2 = (("foo", "bar"), ("a", 1), ("b", 2))
ieq(expect2, table2)
ieq(expect2, table2) # can iterate twice?
def test_records():
"""Test the records function."""
table = (('foo', 'bar'), ('a', 1), ('b', 2))
actual = records(table)
expect = ({'foo': 'a', 'bar': 1}, {'foo': 'b', 'bar': 2})
ieq(expect, actual)
def test_data():
"""Test the data function."""
table = (('foo', 'bar'), ('a', 1), ('b', 2))
actual = data(table)
expect = (('a', 1), ('b', 2))
ieq(expect, actual)
def test_tocsv_appendcsv_gz():
"""Test the tocsv and appendcsv function."""
# exercise function
table = (("foo", "bar"), ("a", 1), ("b", 2), ("c", 2))
f = NamedTemporaryFile(delete=False)
fn = f.name + ".gz"
f.close()
tocsv(table, fn, delimiter="\t")
# check what it did
with gzip.open(fn, "rb") as o:
actual = csv.reader(o, delimiter="\t")
expect = [["foo", "bar"], ["a", "1"], ["b", "2"], ["c", "2"]]
ieq(expect, actual)
# check appending
table2 = (("foo", "bar"), ("d", 7), ("e", 9), ("f", 1))
appendcsv(table2, fn, delimiter="\t")
# check what it did
with gzip.open(fn, "rb") as o:
actual = csv.reader(o, delimiter="\t")
expect = [["foo", "bar"], ["a", "1"], ["b", "2"], ["c", "2"], ["d", "7"], ["e", "9"], ["f", "1"]]
ieq(expect, actual)
def test_operator_overload():
table = (
("foo", "bar", "baz"),
("A", 1, 2),
("B", "2", "3.4"),
("D", "xyz", 9.0),
("B", u"3", u"7.8", True),
("B", "2", 42),
("E", None),
("D", 4, 12.3),
)
f1 = NamedTemporaryFile(delete=False)
p = sort("foo")
p | duplicates("foo") | topickle(f1.name)
p.push(table)
expectation = (
("foo", "bar", "baz"),
("B", "2", "3.4"),
("B", u"3", u"7.8", True),
("B", "2", 42),
("D", "xyz", 9.0),
("D", 4, 12.3),
)
ieq(expectation, frompickle(f1.name))
def test_recordmapmany():
table = (('id', 'sex', 'age', 'height', 'weight'),
(1, 'male', 16, 1.45, 62.0),
(2, 'female', 19, 1.34, 55.4),
(3, '-', 17, 1.78, 74.4),
(4, 'male', 21, 1.33))
def rowgenerator(rec):
transmf = {'male': 'M', 'female': 'F'}
yield [rec['id'], 'gender',
transmf[rec['sex']] if rec['sex'] in transmf else rec['sex']]
yield [rec['id'], 'age_months', rec['age'] * 12]
yield [rec['id'], 'bmi', rec['weight'] / rec['height'] ** 2]
actual = rowmapmany(table, rowgenerator, fields=['subject_id', 'variable',
'value'])
expect = (('subject_id', 'variable', 'value'),
(1, 'gender', 'M'),
(1, 'age_months', 16*12),
(1, 'bmi', 62.0/1.45**2),
(2, 'gender', 'F'),
(2, 'age_months', 19*12),
(2, 'bmi', 55.4/1.34**2),
(3, 'gender', '-'),
(3, 'age_months', 17*12),
(3, 'bmi', 74.4/1.78**2),
(4, 'gender', 'M'),
(4, 'age_months', 21*12))
ieq(expect, actual)
ieq(expect, actual) # can iteratate twice?
def test_sort_3():
table = (("foo", "bar"), ("C", "2"), ("A", "9"), ("A", "6"), ("F", "1"), ("D", "10"))
result = sort(table, "bar")
expectation = (("foo", "bar"), ("F", "1"), ("D", "10"), ("C", "2"), ("A", "6"), ("A", "9"))
ieq(expectation, result)
def test_fold():
t1 = (('id', 'count'), (1, 3), (1, 5), (2, 4), (2, 8))
t2 = fold(t1, 'id', operator.add, 'count', presorted=True)
expect = (('key', 'value'), (1, 8), (2, 12))
ieq(expect, t2)
ieq(expect, t2)
def test_fromxml():
# initial data
f = NamedTemporaryFile(delete=False)
data = """<table>
<tr>
<td>foo</td><td>bar</td>
</tr>
<tr>
<td>a</td><td>1</td>
</tr>
<tr>
<td>b</td><td>2</td>
</tr>
<tr>
<td>c</td><td>2</td>
</tr>
</table>"""
f.write(data)
f.close()
actual = fromxml(f.name, 'tr', 'td')
expect = (('foo', 'bar'),
('a', '1'),
('b', '2'),
('c', '2'))
ieq(expect, actual)
ieq(expect, actual) # verify can iterate twice
def test_rowreduce_empty():
table = (('foo', 'bar'),)
expect = (('foo', 'bar'),)
reducer = lambda key, rows: (key, [r[0] for r in rows])
actual = rowreduce(table, key='foo', reducer=reducer,
fields=('foo', 'bar'))
ieq(expect, actual)
def test_fromxml_5():
# initial data
f = NamedTemporaryFile(delete=False)
data = """<table>
<row>
<foo>a</foo><baz><bar v='1'/><bar v='3'/></baz>
</row>
<row>
<foo>b</foo><baz><bar v='2'/></baz>
</row>
<row>
<foo>c</foo><baz><bar v='2'/></baz>
</row>
</table>"""
f.write(data)
f.close()
actual = fromxml(f.name, 'row', {'foo': 'foo', 'bar': ('baz/bar', 'v')})
expect = (('foo', 'bar'),
('a', ('1', '3')),
('b', '2'),
('c', '2'))
ieq(expect, actual)
ieq(expect, actual) # verify can iterate twice
def test_duplicates():
table = (('foo', 'bar', 'baz'),
('A', 1, 2),
('B', '2', '3.4'),
('D', 'xyz', 9.0),
('B', u'3', u'7.8', True),
('B', '2', 42),
('E', None),
('D', 4, 12.3))
result = duplicates(table, 'foo')
expectation = (('foo', 'bar', 'baz'),
('B', '2', '3.4'),
('B', u'3', u'7.8', True),
('B', '2', 42),
('D', 'xyz', 9.0),
('D', 4, 12.3))
ieq(expectation, result)
# test with compound key
result = duplicates(table, key=('foo', 'bar'))
expectation = (('foo', 'bar', 'baz'),
('B', '2', '3.4'),
('B', '2', 42))
ieq(expectation, result)
def test_fromsqlite3():
"""Test the fromsqlite3 function."""
# initial data
f = NamedTemporaryFile(delete=False)
data = (('a', 1),
('b', 2),
('c', 2.0))
connection = sqlite3.connect(f.name)
c = connection.cursor()
c.execute('create table foobar (foo, bar)')
for row in data:
c.execute('insert into foobar values (?, ?)', row)
connection.commit()
c.close()
connection.close()
# test the function
actual = fromsqlite3(f.name, 'select * from foobar')
expect = (('foo', 'bar'),
('a', 1),
('b', 2),
('c', 2.0))
print list(actual)
ieq(expect, actual, cast=tuple)
ieq(expect, actual, cast=tuple) # verify can iterate twice
def test_sort_none():
table = (('foo', 'bar'),
('C', 2),
('A', 9),
('A', None),
('F', 1),
('D', 10))
result = sort(table, 'bar')
expectation = (('foo', 'bar'),
('A', None),
('F', 1),
('C', 2),
('A', 9),
('D', 10))
ieq(expectation, result)
dt = datetime.now().replace
table = (('foo', 'bar'),
('C', dt(hour=5)),
('A', dt(hour=1)),
('A', None),
('F', dt(hour=9)),
('D', dt(hour=17)))
result = sort(table, 'bar')
expectation = (('foo', 'bar'),
('A', None),
('A', dt(hour=1)),
('C', dt(hour=5)),
('F', dt(hour=9)),
('D', dt(hour=17)))
ieq(expectation, result)
def test_fromdb_mkcursor():
# initial data
data = (('a', 1),
('b', 2),
('c', 2.0))
connection = sqlite3.connect(':memory:')
c = connection.cursor()
c.execute('create table foobar (foo, bar)')
for row in data:
c.execute('insert into foobar values (?, ?)', row)
connection.commit()
c.close()
# test the function
mkcursor = lambda: connection.cursor()
actual = fromdb(mkcursor, 'select * from foobar')
expect = (('foo', 'bar'),
('a', 1),
('b', 2),
('c', 2.0))
ieq(expect, actual)
ieq(expect, actual) # verify can iterate twice
# test iterators are isolated
i1 = iter(actual)
i2 = iter(actual)
eq_(('foo', 'bar'), i1.next())
eq_(('a', 1), i1.next())
eq_(('foo', 'bar'), i2.next())
eq_(('b', 2), i1.next())
def test_sort_2():
table = (('foo', 'bar'),
('C', '2'),
('A', '9'),
('A', '6'),
('F', '1'),
('D', '10'))
result = sort(table, key=('foo', 'bar'))
expectation = (('foo', 'bar'),
('A', '6'),
('A', '9'),
('C', '2'),
('D', '10'),
('F', '1'))
ieq(expectation, result)
result = sort(table) # default is lexical sort
expectation = (('foo', 'bar'),
('A', '6'),
('A', '9'),
('C', '2'),
('D', '10'),
('F', '1'))
ieq(expectation, result)
def test_unflatten_empty():
table1 = (('lines',),)
expect1 = (('f0', 'f1', 'f2'),)
actual1 = unflatten(table1, 'lines', 3)
print list(actual1)
ieq(expect1, actual1)
def test_unique():
table = (
("foo", "bar", "baz"),
("A", 1, 2),
("B", "2", "3.4"),
("D", "xyz", 9.0),
("B", u"3", u"7.8", True),
("B", "2", 42),
("E", None),
("D", 4, 12.3),
)
f1 = NamedTemporaryFile(delete=False)
f2 = NamedTemporaryFile(delete=False)
p = sort("foo")
q = p.pipe(unique("foo"))
q.pipe(topickle(f1.name))
q.pipe("remainder", topickle(f2.name))
p.push(table)
expectation = (("foo", "bar", "baz"), ("A", 1, 2), ("E", None))
ieq(expectation, frompickle(f1.name))
exremainder = (
("foo", "bar", "baz"),
("B", "2", "3.4"),
("B", u"3", u"7.8", True),
("B", "2", 42),
("D", "xyz", 9.0),
("D", 4, 12.3),
)
ieq(exremainder, frompickle(f2.name))
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_topickle_appendpickle():
"""Test the topickle and appendpickle functions."""
# exercise function
table = (("foo", "bar"), ("a", 1), ("b", 2), ("c", 2))
f = NamedTemporaryFile(delete=False)
topickle(table, f.name)
def picklereader(file):
try:
while True:
yield pickle.load(file)
except EOFError:
pass
# check what it did
with open(f.name, "rb") as o:
actual = picklereader(o)
ieq(table, actual)
# check appending
table2 = (("foo", "bar"), ("d", 7), ("e", 9), ("f", 1))
appendpickle(table2, f.name)
# check what it did
with open(f.name, "rb") as o:
actual = picklereader(o)
expect = (("foo", "bar"), ("a", 1), ("b", 2), ("c", 2), ("d", 7), ("e", 9), ("f", 1))
ieq(expect, actual)
