diff --git a/bonobo/examples/environment/env_files/get_passed_env_file.py b/bonobo/examples/environment/env_files/get_passed_env_file.py index bb83e67..bb45540 100644 --- a/bonobo/examples/environment/env_files/get_passed_env_file.py +++ b/bonobo/examples/environment/env_files/get_passed_env_file.py @@ -8,7 +8,9 @@ def extract(): test_user_password = os.getenv('TEST_USER_PASSWORD') path = os.getenv('PATH') - return my_secret, test_user_password, path + yield my_secret + yield test_user_password + yield path def load(s: str): diff --git a/bonobo/examples/environment/env_vars/get_passed_env.py b/bonobo/examples/environment/env_vars/get_passed_env.py index f236ba7..e0c6c45 100644 --- a/bonobo/examples/environment/env_vars/get_passed_env.py +++ b/bonobo/examples/environment/env_vars/get_passed_env.py @@ -8,7 +8,11 @@ def extract(): env_test_number = os.getenv('ENV_TEST_NUMBER', 'number') env_test_string = os.getenv('ENV_TEST_STRING', 'string') env_user = os.getenv('USER') - return env_test_user, env_test_number, env_test_string, env_user + + yield env_test_user + yield env_test_number + yield env_test_string + yield env_user def load(s: str): diff --git a/bonobo/execution/node.py b/bonobo/execution/node.py index 28a20b3..c943ba0 100644 --- a/bonobo/execution/node.py +++ b/bonobo/execution/node.py @@ -1,6 +1,7 @@ import traceback from queue import Empty from time import sleep +from types import GeneratorType from bonobo import settings from bonobo.constants import INHERIT_INPUT, NOT_MODIFIED, BEGIN, END @@ -10,7 +11,6 @@ from bonobo.structs.bags import Bag from bonobo.structs.inputs import Input from bonobo.util import get_name, iserrorbag, isloopbackbag, isdict, istuple from bonobo.util.compat import deprecated_alias -from bonobo.util.iterators import iter_if_not_sequence from bonobo.util.statistics import WithStatistics @@ -120,23 +120,21 @@ class NodeExecutionContext(WithStatistics, LoopingExecutionContext): def handle_results(self, input_bag, results): # self._exec_time += timer.duration # Put data onto output channels - try: - results = iter_if_not_sequence(results) - except TypeError: # not an iterator - if results: - self.send(_resolve(input_bag, results)) - else: - # case with no result, an execution went through anyway, use for stats. - # self._exec_count += 1 - pass - else: - while True: # iterator + + if isinstance(results, GeneratorType): + while True: try: result = next(results) except StopIteration: break else: self.send(_resolve(input_bag, result)) + elif results: + self.send(_resolve(input_bag, results)) + else: + # case with no result, an execution went through anyway, use for stats. + # self._exec_count += 1 + pass def _resolve(input_bag, output): diff --git a/bonobo/util/iterators.py b/bonobo/util/iterators.py index ee45614..1ed09ac 100644 --- a/bonobo/util/iterators.py +++ b/bonobo/util/iterators.py @@ -37,12 +37,3 @@ def tuplize(generator): return tuplized -def iter_if_not_sequence(mixed): - if isinstance(mixed, ( - dict, - list, - str, - bytes, - )): - raise TypeError(type(mixed).__name__) - return iter(mixed) diff --git a/tests/execution/test_node.py b/tests/execution/test_node.py index 23748d4..7870323 100644 --- a/tests/execution/test_node.py +++ b/tests/execution/test_node.py @@ -102,3 +102,26 @@ def test_node_dict_chained(): assert len(output) == 2 assert output[0] == {'id': 1, 'name': 'FOO'} assert output[1] == {'id': 2, 'name': 'BAR'} + +def test_node_tuple(): + def f(): + return 'foo', 'bar' + + with BufferingNodeExecutionContext(f) as context: + context.write_sync(Bag()) + output = context.get_buffer() + + assert len(output) == 1 + assert output[0] == ('foo', 'bar') + + def g(): + yield 'foo', 'bar' + yield 'foo', 'baz' + + with BufferingNodeExecutionContext(g) as context: + context.write_sync(Bag()) + output = context.get_buffer() + + assert len(output) == 2 + assert output[0] == ('foo', 'bar') + assert output[1] == ('foo', 'baz')