def test_inherit():
bag = Bag('a', a=1)
bag2 = Bag.inherit('b', b=2, _parent=bag)
bag3 = bag.extend('c', c=3)
bag4 = Bag('d', d=4)
assert bag.args == ('a', )
assert bag.kwargs == {'a': 1}
assert bag.flags is ()
assert bag2.args == ('a', 'b', )
assert bag2.kwargs == {'a': 1, 'b': 2}
assert INHERIT_INPUT in bag2.flags
assert bag3.args == ('a', 'c', )
assert bag3.kwargs == {'a': 1, 'c': 3}
assert bag3.flags is ()
assert bag4.args == ('d', )
assert bag4.kwargs == {'d': 4}
assert bag4.flags is ()
bag4.set_parent(bag)
assert bag4.args == ('a', 'd', )
assert bag4.kwargs == {'a': 1, 'd': 4}
assert bag4.flags is ()
bag4.set_parent(bag3)
assert bag4.args == ('a', 'c', 'd', )
assert bag4.kwargs == {'a': 1, 'c': 3, 'd': 4}
assert bag4.flags is ()
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 test_pickle():
bag1 = Bag('a', a=1)
bag2 = Bag.inherit('b', b=2, _parent=bag1)
bag3 = bag1.extend('c', c=3)
bag4 = Bag('d', d=4)
# XXX todo this probably won't work with inheriting bags if parent is not there anymore? maybe that's not true
# because the parent may be in the serialization output but we need to verify this assertion.
for bag in bag1, bag2, bag3, bag4:
pickled = pickle.dumps(bag)
unpickled = pickle.loads(pickled)
assert unpickled == bag
def test_write_csv_to_file(tmpdir):
file = tmpdir.join('output.json')
writer = CsvWriter(path=str(file))
context = NodeExecutionContext(writer, None)
context.recv(BEGIN, Bag({'foo': 'bar'}), Bag({'foo': 'baz', 'ignore': 'this'}), END)
context.start()
context.step()
context.step()
context.stop()
assert file.read() == 'foo\nbar\nbaz\n'
with pytest.raises(AttributeError):
getattr(context, 'file')
def test_read_csv_from_file_kwargs(tmpdir):
fs, filename, services = csv_tester.get_services_for_reader(tmpdir)
with CapturingNodeExecutionContext(
CsvReader(path=filename, delimiter=','),
services=services,
) as context:
context.write(BEGIN, Bag(), END)
context.step()
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)
_args, _kwargs = args0[0].get()
assert not len(_args) and _kwargs == {
'a': 'a foo',
'b': 'b foo',
'c': 'c foo',
}
_args, _kwargs = args1[0].get()
assert not len(_args) and _kwargs == {
'a': 'a bar',
'b': 'b bar',
'c': 'c bar',
}
def test_read_csv_from_file(tmpdir):
fs, filename = open_fs(tmpdir), 'input.csv'
fs.open(filename, 'w').write('a,b,c\na foo,b foo,c foo\na bar,b bar,c bar')
reader = CsvReader(path=filename, delimiter=',')
context = CapturingNodeExecutionContext(reader, services={'fs': fs})
context.start()
context.write(BEGIN, Bag(), END)
context.step()
context.stop()
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] == {
'a': 'a foo',
'b': 'b foo',
'c': 'c foo',
}
assert args1[0].args[0] == {
'a': 'a bar',
'b': 'b bar',
'c': 'c bar',
}
def test_read_csv_from_file_arg0(tmpdir):
fs, filename, services = csv_tester.get_services_for_reader(tmpdir)
with CapturingNodeExecutionContext(
CsvReader(path=filename,
delimiter=',',
ioformat=settings.IOFORMAT_ARG0),
services=services,
) as context:
context.write(BEGIN, Bag(), END)
context.step()
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] == {
'a': 'a foo',
'b': 'b foo',
'c': 'c foo',
}
assert args1[0].args[0] == {
'a': 'a bar',
'b': 'b bar',
'c': 'c bar',
}
def test_read_csv_from_file(tmpdir):
file = tmpdir.join('input.csv')
file.write('a,b,c\na foo,b foo,c foo\na bar,b bar,c bar')
reader = CsvReader(path=str(file), delimiter=',')
context = CapturingNodeExecutionContext(reader, None)
context.start()
context.recv(BEGIN, Bag(), END)
context.step()
context.stop()
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] == {
'a': 'a foo',
'b': 'b foo',
'c': 'c foo',
}
assert args1[0].args[0] == {
'a': 'a bar',
'b': 'b bar',
'c': 'c bar',
}
def test_basic():
my_callable1 = Mock()
my_callable2 = Mock()
bag = Bag(*args, **kwargs)
assert not my_callable1.called
result1 = bag.apply(my_callable1)
assert my_callable1.called and result1 is my_callable1.return_value
assert not my_callable2.called
result2 = bag.apply(my_callable2)
assert my_callable2.called and result2 is my_callable2.return_value
assert result1 is not result2
my_callable1.assert_called_once_with(*args, **kwargs)
my_callable2.assert_called_once_with(*args, **kwargs)
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_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_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')
def test_file_reader(tmpdir):
fs, filename, services = txt_tester.get_services_for_reader(tmpdir)
with CapturingNodeExecutionContext(FileReader(path=filename), services=services) as context:
context.write(BEGIN, Bag(), END)
context.step()
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] == 'Hello'
assert args1[0].args[0] == 'World'
def test_write_json_to_file(tmpdir):
file = tmpdir.join('output.json')
writer = JsonWriter(path=str(file))
context = NodeExecutionContext(writer, None)
context.start()
context.recv(BEGIN, Bag({'foo': 'bar'}), END)
context.step()
context.stop()
assert file.read() == '[\n{"foo": "bar"}\n]'
with pytest.raises(AttributeError):
getattr(context, 'file')
with pytest.raises(AttributeError):
getattr(context, 'first')
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_read_json_arg0(tmpdir):
fs, filename, services = json_tester.get_services_for_reader(tmpdir)
with CapturingNodeExecutionContext(
JsonReader(filename, ioformat=settings.IOFORMAT_ARG0),
services=services,
) as context:
context.write(BEGIN, Bag(), END)
context.step()
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'}
def test_simple_execution_context():
graph = Graph()
graph.add_chain(*chain)
ctx = GraphExecutionContext(graph)
assert len(ctx.nodes) == len(chain)
assert not len(ctx.plugins)
for i, node in enumerate(chain):
assert ctx[i].wrapped is node
assert not ctx.alive
ctx.recv(BEGIN, Bag(), END)
assert not ctx.alive
ctx.start()
assert ctx.alive
def test_file_reader_in_context(tmpdir):
file = tmpdir.join('input.txt')
file.write('Hello\nWorld\n')
reader = FileReader(path=str(file))
context = CapturingNodeExecutionContext(reader, None)
context.start()
context.recv(BEGIN, Bag(), END)
context.step()
context.stop()
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] == 'Hello'
assert args1[0].args[0] == 'World'
def test_read_json_from_file(tmpdir):
file = tmpdir.join('input.json')
file.write('[{"x": "foo"},{"x": "bar"}]')
reader = JsonReader(path=str(file))
context = CapturingNodeExecutionContext(reader, None)
context.start()
context.recv(BEGIN, Bag(), END)
context.step()
context.stop()
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'}
def test_read_json_from_file(tmpdir):
fs, filename = open_fs(tmpdir), 'input.json'
fs.open(filename, 'w').write('[{"x": "foo"},{"x": "bar"}]')
reader = JsonReader(path=filename)
context = CapturingNodeExecutionContext(reader, services={'fs': fs})
context.start()
context.recv(BEGIN, Bag(), END)
context.step()
context.stop()
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'}
def test_file_reader_in_context(tmpdir):
fs, filename = open_fs(tmpdir), 'input.txt'
fs.open(filename, 'w').write('Hello\nWorld\n')
reader = FileReader(path=filename)
context = CapturingNodeExecutionContext(reader, services={'fs': fs})
context.start()
context.recv(BEGIN, Bag(), END)
context.step()
context.stop()
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] == 'Hello'
assert args1[0].args[0] == 'World'
def __call__(self, *args, **kwargs):
for operation in self.operations:
args, kwargs = operation.apply(*args, **kwargs)
return Bag(*args, **kwargs)
def test_iterator():
bag = Bag()
assert list(bag.apply([1, 2, 3])) == [1, 2, 3]
assert list(bag.apply((1, 2, 3))) == [1, 2, 3]
assert list(bag.apply(range(5))) == [0, 1, 2, 3, 4]
assert list(bag.apply('azerty')) == ['a', 'z', 'e', 'r', 't', 'y']
def get_fields(**row):
return Bag(**row['fields'])
def __call__(self, *args, **kwargs):
print('factory call on', args, kwargs)
for operation in self.operations:
args, kwargs = operation.apply(*args, **kwargs)
print(' ... after', operation, 'got', args, kwargs)
return Bag(*args, **kwargs)
def transform(topic: str):
return Bag.inherit(title=topic.title(), rand=randint(10, 99))
def extract():
yield Bag(topic='foo')
yield Bag(topic='bar')
yield Bag(topic='baz')
def test_repr():
bag = Bag('a', a=1)
assert repr(bag) == "<Bag ('a', a=1)>"
def test_eq_operator():
assert Bag('foo') == Bag('foo')
assert Bag('foo') != Bag('bar')
assert Bag('foo') is not Bag('foo')
assert Bag('foo') != Token('foo')
assert Token('foo') != Bag('foo')