diff --git a/eliot/_generators.py b/eliot/_generators.py index 1a7b0a1..bb452ef 100644 --- a/eliot/_generators.py +++ b/eliot/_generators.py @@ -36,13 +36,6 @@ def in_generator(self, generator): self._current_generator = previous_generator -class GeneratorSupportNotEnabled(Exception): - """ - An attempt was made to use a decorated generator without first turning on - the generator context manager. - """ - - def eliot_friendly_generator_function(original): """ Decorate a generator function so that the Eliot action context is @@ -50,7 +43,7 @@ def eliot_friendly_generator_function(original): """ @wraps(original) - def wrapper(*a, **kw): + def wrapper(context, *a, **kw): # Keep track of whether the next value to deliver to the generator is # a non-exception or an exception. ok = True @@ -64,8 +57,6 @@ def wrapper(*a, **kw): # generator function can run until we call send or throw on it. gen = original(*a, **kw) - # Initialize the per-generator context to a copy of the current context. - context = copy_context() while True: try: # Whichever way we invoke the generator, we will do it @@ -104,7 +95,7 @@ def go(): # indication of where the yield occurred. # # This is noisy, enable only for debugging: - if wrapper.debug: + if trampoline.debug: log_message(message_type="yielded") return value_out @@ -134,5 +125,12 @@ def go(): else: ok = True - wrapper.debug = False - return wrapper + @wraps(original) + def trampoline(*a, **kw): + # Initialize the generator context to a copy of the current context at + # the site where the generator is invoked. + context = copy_context() + return wrapper(context, *a, **kw) + + trampoline.debug = False + return trampoline diff --git a/eliot/tests/test_generators.py b/eliot/tests/test_generators.py index a92cff3..7a216e9 100644 --- a/eliot/tests/test_generators.py +++ b/eliot/tests/test_generators.py @@ -220,14 +220,14 @@ def g(which): g.debug = True # output yielded messages - gens = [g("1"), g("2")] with start_action(action_type="the-action"): - while gens: - for g in gens[:]: - try: - next(g) - except StopIteration: - gens.remove(g) + gens = [g("1"), g("2")] + while gens: + for g in gens[:]: + try: + next(g) + except StopIteration: + gens.remove(g) assert_expected_action_tree( self, @@ -292,3 +292,39 @@ def g(recurse): } ], ) + + @capture_logging(None) + def test_capture_context(self, logger): + """ + L{eliot_friendly_generator_function} decorated generators capture the + context where they are created, not where L{.send} is first called on + them. + """ + + @eliot_friendly_generator_function + def g(): + yield + + g.debug = True # output yielded messages + + with start_action(action_type="the-action"): + with start_action(action_type="start"): + gen = g() + with start_action(action_type="run"): + list(gen) + + assert_expected_action_tree( + self, + logger, + "the-action", + [ + { + "start": [ + "yielded", + ], + }, + { + "run": [], + }, + ], + )