diff --git a/saucebrush/sources.py b/saucebrush/sources.py index cc1330d..8c161b9 100644 --- a/saucebrush/sources.py +++ b/saucebrush/sources.py @@ -63,7 +63,7 @@ class FixedWidthFileSource(object): line = self._fwfile.next() record = {} for name, range_ in self._fields_dict.iteritems(): - record[name] = line[range[0]:range_[1]].rstrip(self._fillchars) + record[name] = line[range_[0]:range_[1]].rstrip(self._fillchars) return record diff --git a/saucebrush/tests/__init__.py b/saucebrush/tests/__init__.py index f4683f4..d27f9e5 100644 --- a/saucebrush/tests/__init__.py +++ b/saucebrush/tests/__init__.py @@ -1,7 +1,9 @@ import unittest from saucebrush.tests.filters import FilterTestCase +from saucebrush.tests.sources import SourceTestCase filter_suite = unittest.TestLoader().loadTestsFromTestCase(FilterTestCase) +source_suite = unittest.TestLoader().loadTestsFromTestCase(SourceTestCase) if __name__ == '__main__': unittest.main() diff --git a/saucebrush/tests/sources.py b/saucebrush/tests/sources.py new file mode 100644 index 0000000..49c6721 --- /dev/null +++ b/saucebrush/tests/sources.py @@ -0,0 +1,58 @@ +import unittest +import cStringIO +from saucebrush.sources import CSVSource, FixedWidthFileSource + +class SourceTestCase(unittest.TestCase): + + def _get_csv(self): + data = '''a,b,c +1,2,3 +5,5,5 +1,10,100''' + return cStringIO.StringIO(data) + + def test_csv_source_basic(self): + source = CSVSource(self._get_csv()) + expected_data = [{'a':'1', 'b':'2', 'c':'3'}, + {'a':'5', 'b':'5', 'c':'5'}, + {'a':'1', 'b':'10', 'c':'100'}] + self.assertEquals(list(source), expected_data) + + def test_csv_source_fieldnames(self): + source = CSVSource(self._get_csv(), ['x','y','z']) + expected_data = [{'x':'a', 'y':'b', 'z':'c'}, + {'x':'1', 'y':'2', 'z':'3'}, + {'x':'5', 'y':'5', 'z':'5'}, + {'x':'1', 'y':'10', 'z':'100'}] + self.assertEquals(list(source), expected_data) + + def test_csv_source_skiprows(self): + source = CSVSource(self._get_csv(), skiprows=1) + expected_data = [{'a':'5', 'b':'5', 'c':'5'}, + {'a':'1', 'b':'10', 'c':'100'}] + self.assertEquals(list(source), expected_data) + + def test_fixed_width_source(self): + data = cStringIO.StringIO('JamesNovember 3 1986\nTim September151999') + fields = (('name',5), ('month',9), ('day',2), ('year',4)) + source = FixedWidthFileSource(data, fields) + expected_data = [{'name':'James', 'month':'November', 'day':'3', + 'year':'1986'}, + {'name':'Tim', 'month':'September', 'day':'15', + 'year':'1999'}] + self.assertEquals(list(source), expected_data) + + def test_fixed_width_source(self): + data = cStringIO.StringIO('JamesNovember.3.1986\nTim..September151999') + fields = (('name',5), ('month',9), ('day',2), ('year',4)) + source = FixedWidthFileSource(data, fields, fillchars='.') + expected_data = [{'name':'James', 'month':'November', 'day':'3', + 'year':'1986'}, + {'name':'Tim', 'month':'September', 'day':'15', + 'year':'1999'}] + self.assertEquals(list(source), expected_data) + + + +if __name__ == '__main__': + unittest.main()