def test_gen_columns_fun(self):
"""Test column generation with conversion function."""
fn = self.filenames[0]
# Get the first two columns. Both columns are of type int.
def col_fun(fields):
return int(fields[0]), int(fields[1])
data_gen = du.gen_columns(fn, self.header_lines[fn], col_fun)
for i, (c0, c1) in enumerate(data_gen):
self.assertEqual(c0, self.columns[fn][0][i])
self.assertEqual(c1, self.columns[fn][1][i])
# Read columns 1 and 3 and convert them to type
# int and float, respectively. Skip the two header rows.
def col_fun(fields):
return int(fields[1]), float(fields[3])
data_gen = du.gen_columns(fn, self.header_lines[fn], col_fun)
for i, (c1, c3) in enumerate(data_gen):
self.assertEqual(c1, self.columns[fn][1][i])
self.assertEqual(c3, self.columns[fn][3][i])
# Read only the second column. Note the use of tuples!
def col_fun(fields):
return int(fields[1]),
data_gen = du.gen_columns(fn, self.header_lines[fn], col_fun)
for i, (c1,) in enumerate(data_gen):
self.assertEqual(c1, self.columns[fn][1][i])
def test_gen_columns(self):
"""Test basic column generation."""
fn = self.filenames[0]
# Get the first two columns. Both columns are of type int.
data_gen = du.gen_columns(fn, self.header_lines[fn], (int, int))
for i, (c0, c1) in enumerate(data_gen):
self.assertEqual(c0, self.columns[fn][0][i])
self.assertEqual(c1, self.columns[fn][1][i])
# Read columns 1 and 3 and convert them to type
# int and float, respectively. Skip the two header rows.
data_gen = du.gen_columns(fn, self.header_lines[fn], (int, float), (1,3))
for i, (c1, c3) in enumerate(data_gen):
self.assertEqual(c1, self.columns[fn][1][i])
self.assertEqual(c3, self.columns[fn][3][i])
# Read only the second column. Note the use of tuples!
data_gen = du.gen_columns(fn, self.header_lines[fn], (int,), (1,))
for i, (c1,) in enumerate(data_gen):
self.assertEqual(c1, self.columns[fn][1][i])
# Finding a float when int is expected should raise
# ValueError.
fn = self.filenames[2]
data_gen = du.gen_columns(fn, self.header_lines[fn], (int,int), (0,1))
self.assertRaises(ValueError, lambda y: [x for x in y], data_gen)
# A missing column should raise IndexError.
fn = self.filenames[1]
data_gen = du.gen_columns(fn, self.header_lines[fn], (int,float), (1,3))
self.assertRaises(IndexError, lambda y: [x for x in y], data_gen)
