feat: add support for tracing of generators using capture_method decorator

aws_lambda_powertools/tracing/tracer.py

@@ -1,3 +1,4 @@
+import contextlib
 import copy
 import functools
 import inspect
@@ -280,7 +281,7 @@ def decorate(event, context):
 
         return decorate
 
-    def capture_method(self, method: Callable = None):
+    def capture_method(self, method: Callable = None):  # noqa: C901
         """Decorator to create subsegment for arbitrary functions
 
         It also captures both response and exceptions as metadata
@@ -412,6 +413,44 @@ async def decorate(*args, **kwargs):
 
                     return response
 
+        elif inspect.isgeneratorfunction(method):
+
+            @functools.wraps(method)
+            def decorate(*args, **kwargs):
+                with self.provider.in_subsegment(name=f"## {method_name}") as subsegment:
+                    try:
+                        logger.debug(f"Calling method: {method_name}")
+                        result = yield from method(*args, **kwargs)
+                        self._add_response_as_metadata(function_name=method_name, data=result, subsegment=subsegment)
+                    except Exception as err:
+                        logger.exception(f"Exception received from '{method_name}' method")
+                        self._add_full_exception_as_metadata(
+                            function_name=method_name, error=err, subsegment=subsegment
+                        )
+                        raise
+
+                    return result
+
+        elif hasattr(method, "__wrapped__") and inspect.isgeneratorfunction(method.__wrapped__):
+
+            @functools.wraps(method)
+            @contextlib.contextmanager
+            def decorate(*args, **kwargs):
+                with self.provider.in_subsegment(name=f"## {method_name}") as subsegment:
+                    try:
+                        logger.debug(f"Calling method: {method_name}")
+                        with method(*args, **kwargs) as return_val:
+                            result = return_val
+                        self._add_response_as_metadata(function_name=method_name, data=result, subsegment=subsegment)
+                    except Exception as err:
+                        logger.exception(f"Exception received from '{method_name}' method")
+                        self._add_full_exception_as_metadata(
+                            function_name=method_name, error=err, subsegment=subsegment
+                        )
+                        raise
+
+                    yield result
+
         else:
 
             @functools.wraps(method)

example/tests/test_handler.py

@@ -4,6 +4,14 @@
 
 import pytest
 
+from aws_lambda_powertools import Tracer
+
+
+@pytest.fixture(scope="function", autouse=True)
+def reset_tracing_config():
+    Tracer._reset_config()
+    yield
+
 
 @pytest.fixture()
 def env_vars(monkeypatch):

tests/functional/test_tracing.py

@@ -1,3 +1,5 @@
+import contextlib
+
 import pytest
 
 from aws_lambda_powertools import Tracer
@@ -150,3 +152,36 @@ def sums_values():
         return func_1() + func_2()
 
     sums_values()
+
+
+def test_tracer_yield_with_capture():
+    # GIVEN tracer method decorator is used
+    tracer = Tracer(disabled=True)
+
+    # WHEN capture_method decorator is applied to a context manager
+    @tracer.capture_method
+    @contextlib.contextmanager
+    def yield_with_capture():
+        yield "testresult"
+
+    # Or WHEN capture_method decorator is applied to a generator function
+    @tracer.capture_method
+    def generator_func():
+        yield "testresult2"
+
+    @tracer.capture_lambda_handler
+    def handler(event, context):
+        result = []
+        with yield_with_capture() as yielded_value:
+            result.append(yielded_value)
+
+        gen = generator_func()
+
+        result.append(next(gen))
+
+        return result
+
+    # THEN no exception is thrown, and the functions properly return values
+    result = handler({}, {})
+    assert "testresult" in result
+    assert "testresult2" in result

tests/unit/test_tracing.py

@@ -1,3 +1,4 @@
+import contextlib
 import sys
 from typing import NamedTuple
 from unittest import mock
@@ -348,3 +349,65 @@ async def greeting(name, message):
     put_metadata_mock_args = in_subsegment_mock.put_metadata.call_args[1]
     assert put_metadata_mock_args["key"] == "greeting error"
     assert put_metadata_mock_args["namespace"] == "booking"
+
+
+def test_tracer_yield_from_context_manager(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN capture_method decorator is used on a context manager
+    @tracer.capture_method
+    @contextlib.contextmanager
+    def yield_with_capture():
+        yield "test result"
+
+    @tracer.capture_lambda_handler
+    def handler(event, context):
+        response = []
+        with yield_with_capture() as yielded_value:
+            response.append(yielded_value)
+
+        return response
+
+    result = handler({}, {})
+
+    # THEN we should have a subsegment named after the method name
+    # and add its response as trace metadata
+    handler_trace, yield_function_trace = in_subsegment_mock.in_subsegment.call_args_list
+
+    assert "test result" in in_subsegment_mock.put_metadata.call_args[1]["value"]
+    assert in_subsegment_mock.in_subsegment.call_count == 2
+    assert handler_trace == mocker.call(name="## handler")
+    assert yield_function_trace == mocker.call(name="## yield_with_capture")
+    assert "test result" in result
+
+
+def test_tracer_yield_from_generator(mocker, provider_stub, in_subsegment_mock):
+    # GIVEN tracer is initialized
+    provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+    tracer = Tracer(provider=provider, service="booking")
+
+    # WHEN capture_method decorator is used on a generator function
+    @tracer.capture_method
+    def generator_fn():
+        yield "test result"
+
+    @tracer.capture_lambda_handler
+    def handler(event, context):
+        gen = generator_fn()
+        response = list(gen)
+
+        return response
+
+    result = handler({}, {})
+
+    # THEN we should have a subsegment named after the method name
+    # and add its response as trace metadata
+    handler_trace, generator_fn_trace = in_subsegment_mock.in_subsegment.call_args_list
+
+    assert "test result" in in_subsegment_mock.put_metadata.call_args[1]["value"]
+    assert in_subsegment_mock.in_subsegment.call_count == 2
+    assert handler_trace == mocker.call(name="## handler")
+    assert generator_fn_trace == mocker.call(name="## generator_fn")
+    assert "test result" in result