diff --git a/bonobo/nodes/io/csv.py b/bonobo/nodes/io/csv.py index c3a9b5d..6fd0c5e 100644 --- a/bonobo/nodes/io/csv.py +++ b/bonobo/nodes/io/csv.py @@ -54,16 +54,10 @@ class CsvHandler(FileHandler): @use_context class CsvReader(FileReader, CsvHandler): """ - Reads a CSV and yield the values as dicts. + Reads a CSV and yield the values. """ - skip = Option( - int, - default=0, - __doc__=""" - If set and greater than zero, the reader will skip this amount of lines. - """, - ) + skip = Option(int, default=0, __doc__="If set and greater than zero, the reader will skip this amount of lines.") @Method( positional=False, @@ -96,6 +90,8 @@ class CsvReader(FileReader, CsvHandler): @use_context class CsvWriter(FileWriter, CsvHandler): + skip_header = Option(bool, default=False, __doc__="If true, the writer will not produce a file header.") + @Method( __doc__=""" Builds the CSV writer, a.k.a an object we can pass a field collection to be written as one line in the @@ -114,7 +110,7 @@ class CsvWriter(FileWriter, CsvHandler): if not context.lineno: context.writer = self.writer_factory(file) - if fields: + if fields and not self.skip_header: context.writer(fields) context.lineno += 1 diff --git a/tests/nodes/io/test_csv.py b/tests/nodes/io/test_csv.py index 131c88e..6fda2e9 100644 --- a/tests/nodes/io/test_csv.py +++ b/tests/nodes/io/test_csv.py @@ -89,6 +89,22 @@ class CsvWriterTest(Csv, WriterTest, TestCase): assert self.readlines() == ("foo,bar", "a,b", "c,d") + @incontext(skip_header=False) + def test_fields_with_headers(self, context): + context.set_input_fields(["foo", "bar"]) + context.write_sync(("a", "b"), ("c", "d")) + context.stop() + + assert self.readlines() == ("foo,bar", "a,b", "c,d") + + @incontext(skip_header=True) + def test_fields_without_headers(self, context): + context.set_input_fields(["foo", "bar"]) + context.write_sync(("a", "b"), ("c", "d")) + context.stop() + + assert self.readlines() == ("a,b", "c,d") + @incontext() def test_fields_from_type(self, context): context.set_input_type(namedtuple("Point", "x y"))