From 689bd4332eef1e88e1321ef6b4168bb532c8aa46 Mon Sep 17 00:00:00 2001 From: Blazej Gorny Date: Wed, 3 Dec 2025 14:56:08 +0100 Subject: [PATCH 1/2] fix: correct async generator handling in @metric_scope decorator - Replace manual __anext__() calls with async for loop syntax - Replace manual next() calls with for loop syntax - Change exception handling from catching StopIteration to using try-finally - Async generators now properly complete without raising exceptions - Ensure metrics are always flushed even when exceptions occur Fixes #128 --- aws_embedded_metrics/metric_scope/__init__.py | 16 +++------ tests/metric_scope/test_metric_scope.py | 35 +++++++++++++++++-- 2 files changed, 37 insertions(+), 14 deletions(-) diff --git a/aws_embedded_metrics/metric_scope/__init__.py b/aws_embedded_metrics/metric_scope/__init__.py index 3389945..9c90c18 100644 --- a/aws_embedded_metrics/metric_scope/__init__.py +++ b/aws_embedded_metrics/metric_scope/__init__.py @@ -29,15 +29,11 @@ async def async_gen_wrapper(*args, **kwargs): # type: ignore kwargs["metrics"] = logger try: - fn_gen = fn(*args, **kwargs) - while True: - result = await fn_gen.__anext__() + async for result in fn(*args, **kwargs): await logger.flush() yield result - except Exception as ex: + finally: await logger.flush() - if not isinstance(ex, StopIteration): - raise return cast(F, async_gen_wrapper) @@ -49,15 +45,11 @@ def gen_wrapper(*args, **kwargs): # type: ignore kwargs["metrics"] = logger try: - fn_gen = fn(*args, **kwargs) - while True: - result = next(fn_gen) + for result in fn(*args, **kwargs): asyncio.run(logger.flush()) yield result - except Exception as ex: + finally: asyncio.run(logger.flush()) - if not isinstance(ex, StopIteration): - raise return cast(F, gen_wrapper) diff --git a/tests/metric_scope/test_metric_scope.py b/tests/metric_scope/test_metric_scope.py index 9ebd1f1..72c0443 100644 --- a/tests/metric_scope/test_metric_scope.py +++ b/tests/metric_scope/test_metric_scope.py @@ -169,7 +169,38 @@ def my_handler(metrics): assert expected_timestamp_second == actual_timestamp_second -def test_sync_scope_iterates_generator(mock_logger): +@pytest.mark.asyncio +async def test_async_generator_completes_successfully(mock_logger): + expected_results = [1, 2, 3] + + @metric_scope + async def my_handler(): + for item in expected_results: + yield item + + actual_results = [] + async for result in my_handler(): + actual_results.append(result) + + assert actual_results == expected_results + assert InvocationTracker.invocations == 4 # 3 yields + 1 final flush + + +def test_sync_generator_completes_successfully(mock_logger): + expected_results = [1, 2, 3] + + @metric_scope + def my_handler(): + yield from expected_results + + actual_results = [] + for result in my_handler(): + actual_results.append(result) + + assert actual_results == expected_results + assert InvocationTracker.invocations == 4 # 3 yields + 1 final flush + +def test_sync_generator_handles_exception(mock_logger): expected_results = [1, 2] @metric_scope @@ -187,7 +218,7 @@ def my_handler(): @pytest.mark.asyncio -async def test_async_scope_iterates_async_generator(mock_logger): +async def test_async_generator_handles_exception(mock_logger): expected_results = [1, 2] @metric_scope From 8388419506f972cc0a121a302d4556609cb3c9d9 Mon Sep 17 00:00:00 2001 From: Blazej Gorny Date: Wed, 3 Dec 2025 14:59:54 +0100 Subject: [PATCH 2/2] perf: flush metrics once at generator completion instead of after each yield - Remove flush calls after each yield in generators - Metrics are now flushed only when generator completes (in finally block) - Improves performance by reducing flush operations - Maintains correctness with guaranteed flush on completion or exception --- aws_embedded_metrics/metric_scope/__init__.py | 2 -- tests/metric_scope/test_metric_scope.py | 8 ++++---- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/aws_embedded_metrics/metric_scope/__init__.py b/aws_embedded_metrics/metric_scope/__init__.py index 9c90c18..4d0ff34 100644 --- a/aws_embedded_metrics/metric_scope/__init__.py +++ b/aws_embedded_metrics/metric_scope/__init__.py @@ -30,7 +30,6 @@ async def async_gen_wrapper(*args, **kwargs): # type: ignore try: async for result in fn(*args, **kwargs): - await logger.flush() yield result finally: await logger.flush() @@ -46,7 +45,6 @@ def gen_wrapper(*args, **kwargs): # type: ignore try: for result in fn(*args, **kwargs): - asyncio.run(logger.flush()) yield result finally: asyncio.run(logger.flush()) diff --git a/tests/metric_scope/test_metric_scope.py b/tests/metric_scope/test_metric_scope.py index 72c0443..31fe4d2 100644 --- a/tests/metric_scope/test_metric_scope.py +++ b/tests/metric_scope/test_metric_scope.py @@ -183,7 +183,7 @@ async def my_handler(): actual_results.append(result) assert actual_results == expected_results - assert InvocationTracker.invocations == 4 # 3 yields + 1 final flush + assert InvocationTracker.invocations == 1 def test_sync_generator_completes_successfully(mock_logger): @@ -198,7 +198,7 @@ def my_handler(): actual_results.append(result) assert actual_results == expected_results - assert InvocationTracker.invocations == 4 # 3 yields + 1 final flush + assert InvocationTracker.invocations == 1 def test_sync_generator_handles_exception(mock_logger): expected_results = [1, 2] @@ -214,7 +214,7 @@ def my_handler(): actual_results.append(result) assert actual_results == expected_results - assert InvocationTracker.invocations == 3 + assert InvocationTracker.invocations == 1 @pytest.mark.asyncio @@ -234,7 +234,7 @@ async def my_handler(): actual_results.append(result) assert actual_results == expected_results - assert InvocationTracker.invocations == 3 + assert InvocationTracker.invocations == 1 # Test helpers