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)