Allows to provide False fields to CsvWriter, or to override field headers using the fields= option.

This commit is contained in:
Romain Dorgueil
2018-10-27 15:29:19 +02:00
parent a140e6506e
commit 5499c548b0
3 changed files with 22 additions and 3 deletions

View File

@ -6,6 +6,7 @@ from bonobo.constants import NOT_MODIFIED
from bonobo.nodes.io.base import FileHandler
from bonobo.nodes.io.file import FileReader, FileWriter
from bonobo.util import ensure_tuple
from bonobo.util.collections import tuple_or_const
class CsvHandler(FileHandler):
@ -36,7 +37,7 @@ class CsvHandler(FileHandler):
# Fields (renamed from headers)
headers = RenamedOption("fields")
fields = Option(ensure_tuple, required=False)
fields = Option(tuple_or_const, required=False)
def get_dialect_kwargs(self):
return {
@ -108,7 +109,7 @@ class CsvWriter(FileWriter, CsvHandler):
def write(self, file, context, *values, fs):
context.setdefault("lineno", 0)
fields = context.get_input_fields()
fields = context.get_input_fields() if self.fields is None else self.fields
if not context.lineno:
context.writer = self.writer_factory(file)

View File

@ -1,5 +1,6 @@
import bisect
import functools
from collections import Sequence
class sortedlist(list):
@ -32,6 +33,16 @@ def _with_length_check(f):
return _wrapped
def tuple_or_const(tuple_or_mixed, *, consts=(None, False), **kwargs):
if tuple_or_mixed in consts:
return tuple_or_mixed
if isinstance(tuple_or_mixed, str):
pass
elif isinstance(tuple_or_mixed, Sequence):
tuple_or_mixed = tuple(tuple_or_mixed)
return ensure_tuple(tuple_or_mixed, **kwargs)
@_with_length_check
def ensure_tuple(tuple_or_mixed, *, cls=None):
"""