Files
bonobo/bonobo/io/file.py

101 lines
2.4 KiB
Python

from bonobo.config import Configurable, Option
from bonobo.context import ContextProcessor
from bonobo.context.processors import contextual
from bonobo.util.objects import ValueHolder
__all__ = [
'FileReader',
'FileWriter',
]
@contextual
class FileHandler(Configurable):
"""
Abstract component factory for file-related components.
"""
path = Option(str, required=True)
eol = Option(str, default='\n')
mode = Option(str)
@ContextProcessor
def file(self, context):
with self.open() as file:
yield file
def open(self):
return open(self.path, self.mode)
class Reader(FileHandler):
def __call__(self, *args):
yield from self.read(*args)
def read(self, *args):
raise NotImplementedError('Abstract.')
class Writer(FileHandler):
def __call__(self, *args):
return self.write(*args)
def write(self, *args):
raise NotImplementedError('Abstract.')
class FileReader(Reader):
"""
Component factory for file-like readers.
On its own, it can be used to read a file and yield one row per line, trimming the "eol" character at the end if
present. Extending it is usually the right way to create more specific file readers (like json, csv, etc.)
"""
mode = Option(str, default='r')
def read(self, file):
"""
Write a row on the next line of given file.
Prefix is used for newlines.
:param ctx:
:param row:
"""
for line in file:
yield line.rstrip(self.eol)
@contextual
class FileWriter(Writer):
"""
Component factory for file or file-like writers.
On its own, it can be used to write in a file one line per row that comes into this component. Extending it is
usually the right way to create more specific file writers (like json, csv, etc.)
"""
mode = Option(str, default='w+')
@ContextProcessor
def lineno(self, context, file):
lineno = ValueHolder(0, type=int)
yield lineno
def write(self, file, lineno, row):
"""
Write a row on the next line of opened file in context.
:param file fp:
:param str row:
:param str prefix:
"""
self._write_line(file, (self.eol if lineno.value else '') + row)
lineno.value += 1
def _write_line(self, file, line):
return file.write(line)