diff --git a/bonobo/__init__.py b/bonobo/__init__.py index 03ca92a..6832651 100644 --- a/bonobo/__init__.py +++ b/bonobo/__init__.py @@ -23,7 +23,7 @@ import os import sys from .core import * -from .io import * +from .io import FileReader, FileWriter, JsonReader, JsonWriter from .util import * PY35 = (sys.version_info >= (3, 5)) @@ -36,8 +36,10 @@ with open(os.path.realpath(os.path.join(os.path.dirname(__file__), '../version.t __all__ = [ 'Bag', + 'FileReader', 'FileWriter', 'Graph', + 'JsonReader', 'JsonWriter', 'NOT_MODIFIED', 'NaiveStrategy', diff --git a/bonobo/io/__init__.py b/bonobo/io/__init__.py index 983900c..36addbc 100644 --- a/bonobo/io/__init__.py +++ b/bonobo/io/__init__.py @@ -1,11 +1,11 @@ """ Readers and writers for common file formats. """ -from .file import Handler, FileReader, FileWriter -from .json import JsonWriter +from .file import FileReader, FileWriter +from .json import JsonReader, JsonWriter __all__ = [ - 'Handler', 'FileReader', 'FileWriter', + 'JsonReader', 'JsonWriter', ] diff --git a/bonobo/io/file.py b/bonobo/io/file.py index 4d28d03..0b341a2 100644 --- a/bonobo/io/file.py +++ b/bonobo/io/file.py @@ -1,14 +1,14 @@ from bonobo.util.lifecycle import with_context __all__ = [ - 'Handler', + 'FileHandler', 'FileReader', 'FileWriter', ] @with_context -class Handler: +class FileHandler: """ Abstract component factory for file-related components. @@ -44,7 +44,7 @@ class Handler: del ctx.file -class Reader(Handler): +class Reader(FileHandler): def __call__(self, ctx): yield from self.handle(ctx) @@ -52,7 +52,7 @@ class Reader(Handler): raise NotImplementedError('Abstract.') -class Writer(Handler): +class Writer(FileHandler): def __call__(self, ctx, row): return self.handle(ctx, row) diff --git a/bonobo/io/json.py b/bonobo/io/json.py index 120f6c4..b12a811 100644 --- a/bonobo/io/json.py +++ b/bonobo/io/json.py @@ -1,6 +1,6 @@ import json -from .file import FileWriter +from .file import FileWriter, FileReader __all__ = ['JsonWriter', ] @@ -9,9 +9,14 @@ class JsonHandler: eol = ',\n' +class JsonReader(JsonHandler, FileReader): + def handle(self, ctx): + for line in json.load(ctx.file): + yield line + + class JsonWriter(JsonHandler, FileWriter): def initialize(self, ctx): - print('EOL', self.eol) super().initialize(ctx) ctx.file.write('[\n') diff --git a/bonobo/util/testing.py b/bonobo/util/testing.py new file mode 100644 index 0000000..30431c3 --- /dev/null +++ b/bonobo/util/testing.py @@ -0,0 +1,9 @@ +from unittest.mock import MagicMock + +from bonobo.core.contexts import ComponentExecutionContext + + +class CapturingComponentExecutionContext(ComponentExecutionContext): + def __init__(self, component, parent): + super().__init__(component, parent) + self.send = MagicMock() diff --git a/tests/io/test_file.py b/tests/io/test_file.py index 34e7119..32566a5 100644 --- a/tests/io/test_file.py +++ b/tests/io/test_file.py @@ -1,15 +1,11 @@ import pytest -from mock import MagicMock from bonobo import FileWriter, Bag, FileReader from bonobo.core.contexts import ComponentExecutionContext +from bonobo.util.testing import CapturingComponentExecutionContext from bonobo.util.tokens import BEGIN, END -class CapturingComponentExecutionContext(ComponentExecutionContext): - send = MagicMock() - - @pytest.mark.parametrize( 'lines,output', [ diff --git a/tests/io/test_json.py b/tests/io/test_json.py index fd0bea5..b61c082 100644 --- a/tests/io/test_json.py +++ b/tests/io/test_json.py @@ -1,7 +1,8 @@ import pytest -from bonobo import Bag, JsonWriter +from bonobo import Bag, JsonWriter, JsonReader from bonobo.core.contexts import ComponentExecutionContext +from bonobo.util.testing import CapturingComponentExecutionContext from bonobo.util.tokens import BEGIN, END @@ -33,3 +34,26 @@ def test_write_json_without_initializer_should_not_work(tmpdir): context = ComponentExecutionContext(json_writer, None) with pytest.raises(AttributeError): json_writer(context, {'foo': 'bar'}) + + +def test_read_json_from_file(tmpdir): + file = tmpdir.join('input.json') + file.write('[{"x": "foo"},{"x": "bar"}]') + reader = JsonReader(str(file)) + + context = CapturingComponentExecutionContext(reader, None) + + context.initialize() + context.recv(BEGIN, Bag(), END) + context.step() + context.finalize() + + assert len(context.send.mock_calls) == 2 + + args0, kwargs0 = context.send.call_args_list[0] + assert len(args0) == 1 and not len(kwargs0) + args1, kwargs1 = context.send.call_args_list[1] + assert len(args1) == 1 and not len(kwargs1) + + assert args0[0].args[0] == {'x': 'foo'} + assert args1[0].args[0] == {'x': 'bar'}