def test_file_writer_in_context(tmpdir, lines, output):
fs, filename = open_fs(tmpdir), 'output.txt'
writer = FileWriter(path=filename)
context = NodeExecutionContext(writer, services={'fs': fs})
context.start()
context.recv(BEGIN, *map(Bag, lines), END)
for _ in range(len(lines)):
context.step()
context.stop()
assert fs.open(filename).read() == output
def __init__(self, graph, plugins=None, services=None):
self.graph = graph
self.nodes = [
NodeExecutionContext(node, parent=self) for node in self.graph
]
self.plugins = [
PluginExecutionContext(plugin, parent=self)
for plugin in plugins or ()
]
self.services = create_container(services)
# Probably not a good idea to use it unless you really know what you're doing. But you can access the context.
self.services['__graph_context'] = self
for i, node_context in enumerate(self):
node_context.outputs = [
self[j].input for j in self.graph.outputs_of(i)
]
node_context.input.on_begin = partial(node_context.send,
BEGIN,
_control=True)
node_context.input.on_end = partial(node_context.send,
END,
_control=True)
node_context.input.on_finalize = partial(node_context.stop)
def test_write_csv_to_file(tmpdir):
fs, filename = open_fs(tmpdir), 'output.csv'
writer = CsvWriter(path=filename)
context = NodeExecutionContext(writer, services={'fs': fs})
context.write(BEGIN, Bag({'foo': 'bar'}), Bag({'foo': 'baz', 'ignore': 'this'}), END)
context.start()
context.step()
context.step()
context.stop()
assert fs.open(filename).read() == 'foo\nbar\nbaz\n'
with pytest.raises(AttributeError):
getattr(context, 'file')
def test_write_json_to_file(tmpdir):
fs, filename = open_fs(tmpdir), 'output.json'
writer = JsonWriter(path=filename)
context = NodeExecutionContext(writer, services={'fs': fs})
context.start()
context.recv(BEGIN, Bag({'foo': 'bar'}), END)
context.step()
context.stop()
assert fs.open(filename).read() == '[{"foo": "bar"}]'
with pytest.raises(AttributeError):
getattr(context, 'file')
with pytest.raises(AttributeError):
getattr(context, 'first')
def test_write_json_kwargs(tmpdir, add_kwargs):
fs, filename, services = json_tester.get_services_for_writer(tmpdir)
with NodeExecutionContext(JsonWriter(filename, **add_kwargs), services=services) as context:
context.write(BEGIN, Bag(**{'foo': 'bar'}), END)
context.step()
with fs.open(filename) as fp:
assert fp.read() == '[{"foo": "bar"}]'
def test_write_json_arg0(tmpdir):
fs, filename, services = json_tester.get_services_for_writer(tmpdir)
with NodeExecutionContext(JsonWriter(filename, ioformat=settings.IOFORMAT_ARG0), services=services) as context:
context.write(BEGIN, Bag({'foo': 'bar'}), END)
context.step()
with fs.open(filename) as fp:
assert fp.read() == '[{"foo": "bar"}]'
def test_file_writer_in_context(tmpdir, lines, output):
fs, filename, services = txt_tester.get_services_for_writer(tmpdir)
with NodeExecutionContext(FileWriter(path=filename), services=services) as context:
context.write(BEGIN, *map(Bag, lines), END)
for _ in range(len(lines)):
context.step()
with fs.open(filename) as fp:
assert fp.read() == output
def test_write_csv_to_file_kwargs(tmpdir, add_kwargs):
fs, filename, services = csv_tester.get_services_for_writer(tmpdir)
with NodeExecutionContext(CsvWriter(path=filename, **add_kwargs), services=services) as context:
context.write(BEGIN, Bag(**{'foo': 'bar'}), Bag(**{'foo': 'baz', 'ignore': 'this'}), END)
context.step()
context.step()
with fs.open(filename) as fp:
assert fp.read() == 'foo\nbar\nbaz\n'
with pytest.raises(AttributeError):
getattr(context, 'file')
def __init__(self, graph, plugins=None, services=None):
self.graph = graph
self.nodes = [NodeExecutionContext(node, parent=self) for node in self.graph.nodes]
self.plugins = [PluginExecutionContext(plugin, parent=self) for plugin in plugins or ()]
self.services = Container(services) if services else Container()
for i, node_context in enumerate(self):
try:
node_context.outputs = [self[j].input for j in self.graph.outputs_of(i)]
except KeyError:
continue
node_context.input.on_begin = partial(node_context.send, BEGIN, _control=True)
node_context.input.on_end = partial(node_context.send, END, _control=True)
node_context.input.on_finalize = partial(node_context.stop)
def test_write_pickled_dict_to_file(tmpdir):
fs, filename, services = pickle_tester.get_services_for_writer(tmpdir)
with NodeExecutionContext(PickleWriter(filename,
ioformat=settings.IOFORMAT_ARG0),
services=services) as context:
context.write(BEGIN, Bag({'foo': 'bar'}),
Bag({
'foo': 'baz',
'ignore': 'this'
}), END)
context.step()
context.step()
with fs.open(filename, 'rb') as fp:
assert pickle.loads(fp.read()) == {'foo': 'bar'}
with pytest.raises(AttributeError):
getattr(context, 'file')