Capture contect from first call in eliot_friendly_generator_function.

Author: tomprince

eliot/_generators.py

@@ -43,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
@@ -57,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
@@ -97,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
 
@@ -127,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

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,38 @@ 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": [],
+                },
+            ],
+        )