Merge pull request #331 from hartym/302_csv_skip_headers_option

CSV skip headers option
This commit is contained in:
Romain Dorgueil
2019-06-01 09:55:29 +02:00
committed by GitHub
2 changed files with 21 additions and 9 deletions

View File

@ -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

View File

@ -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"))