diff --git a/bonobo/io/__init__.py b/bonobo/io/__init__.py index 22229d9..983900c 100644 --- a/bonobo/io/__init__.py +++ b/bonobo/io/__init__.py @@ -1,10 +1,10 @@ """ Readers and writers for common file formats. """ -from .file import FileHandler, FileReader, FileWriter +from .file import Handler, FileReader, FileWriter from .json import JsonWriter __all__ = [ - 'FileHandler', + 'Handler', 'FileReader', 'FileWriter', 'JsonWriter', diff --git a/bonobo/io/file.py b/bonobo/io/file.py index fba076a..4d28d03 100644 --- a/bonobo/io/file.py +++ b/bonobo/io/file.py @@ -1,26 +1,25 @@ -from functools import partial - from bonobo.util.lifecycle import with_context __all__ = [ - 'FileHandler', + 'Handler', 'FileReader', 'FileWriter', ] @with_context -class FileHandler: +class Handler: """ Abstract component factory for file-related components. """ + eol = '\n' mode = None - def __init__(self, path_or_buf, eol='\n'): + def __init__(self, path_or_buf, eol=None): self.path_or_buf = path_or_buf - self.eol = eol + self.eol = eol or self.eol def open(self): return open(self.path_or_buf, self.mode) @@ -40,26 +39,28 @@ class FileHandler: assert not hasattr(ctx, 'file'), 'A file pointer is already in the context... I do not know what to say...' ctx.file = self.open() - def __call__(self, ctx, *args): - """ - :param ComponentExecutionContext ctx: - :param mixed row: - """ - result = self.handle(ctx, *args) - try: - yield from result - except TypeError: - return - - def handle(self, ctx, *args): - raise NotImplementedError('Abstract.') - def finalize(self, ctx): self.close(ctx.file) del ctx.file -class FileReader(FileHandler): +class Reader(Handler): + def __call__(self, ctx): + yield from self.handle(ctx) + + def handle(self, ctx): + raise NotImplementedError('Abstract.') + + +class Writer(Handler): + def __call__(self, ctx, row): + return self.handle(ctx, row) + + def handle(self, ctx, row): + raise NotImplementedError('Abstract.') + + +class FileReader(Reader): """ Component factory for file-like readers. @@ -70,7 +71,7 @@ class FileReader(FileHandler): mode = 'r' - def handle(self, ctx, *args): + def handle(self, ctx): """ Write a row on the next line of file pointed by `ctx.file`. Prefix is used for newlines. @@ -78,12 +79,11 @@ class FileReader(FileHandler): :param ctx: :param row: """ - assert not len(args) for line in ctx.file: yield line.rstrip(self.eol) -class FileWriter(FileHandler): +class FileWriter(Writer): """ Component factory for file or file-like writers. @@ -100,13 +100,12 @@ class FileWriter(FileHandler): def handle(self, ctx, row): """ - Write a row on the next line of file pointed by fp. Prefix is used for newlines. + Write a row on the next line of opened file in context. :param file fp: :param str row: :param str prefix: """ - self.write(ctx.file, (self.eol if ctx.line else '') + row) ctx.line += 1 diff --git a/bonobo/io/json.py b/bonobo/io/json.py index 156d385..120f6c4 100644 --- a/bonobo/io/json.py +++ b/bonobo/io/json.py @@ -1,23 +1,29 @@ import json from .file import FileWriter -from bonobo.util.lifecycle import with_context __all__ = ['JsonWriter', ] -@with_context -class JsonWriter(FileWriter): - def __init__(self, path_or_buf): - super().__init__(path_or_buf, eol=',\n') +class JsonHandler: + eol = ',\n' + +class JsonWriter(JsonHandler, FileWriter): def initialize(self, ctx): + print('EOL', self.eol) super().initialize(ctx) - ctx.fp.write('[\n') + ctx.file.write('[\n') - def write(self, fp, line, prefix=''): - fp.write(prefix + json.dumps(line)) + def handle(self, ctx, row): + """ + Write a json row on the next line of file pointed by ctx.file. + + :param ctx: + :param row: + """ + return super().handle(ctx, json.dumps(row)) def finalize(self, ctx): - ctx.fp.write('\n]') + ctx.file.write('\n]') super().finalize(ctx) diff --git a/tests/io/test_json.py b/tests/io/test_json.py index 547f435..fd0bea5 100644 --- a/tests/io/test_json.py +++ b/tests/io/test_json.py @@ -20,7 +20,7 @@ def test_write_json_to_file(tmpdir): ]''' with pytest.raises(AttributeError): - getattr(context, 'fp') + getattr(context, 'file') with pytest.raises(AttributeError): getattr(context, 'first')