From 81888a369cf78575c1f8c13f407c7e78eab93ffd Mon Sep 17 00:00:00 2001 From: Frank Hoffmann <15r10nk-git@polarbit.de> Date: Sun, 24 Sep 2023 16:00:08 +0200 Subject: [PATCH] feat: support for __next__ --- executing/_position_node_finder.py | 2 +- tests/test_main.py | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/executing/_position_node_finder.py b/executing/_position_node_finder.py index ca3354c..8ca21a6 100644 --- a/executing/_position_node_finder.py +++ b/executing/_position_node_finder.py @@ -593,7 +593,7 @@ def node_match(node_type: Union[Type, Tuple[Type, ...]], **kwargs: Any) -> bool: if node_match( (ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp, ast.For) - ) and inst_match("GET_ITER"): + ) and inst_match(("GET_ITER", "FOR_ITER")): return if sys.version_info >= (3, 12): diff --git a/tests/test_main.py b/tests/test_main.py index c30628d..7e33247 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -581,12 +581,18 @@ def test_iter(self): class iter_test: def __init__(self, typ): self.typ = typ + self.it = iter([1, 2]) def __iter__(self): assert isinstance(calling_expression(), self.typ) - return iter([1, 2]) + return self - iter(iter_test(ast.Call)) + def __next__(self): + assert isinstance(calling_expression(), self.typ) + return next(self.it) + + assert list(iter_test(ast.Call)) == [1, 2] + assert next(iter(iter_test(ast.Call))) == 1 if sys.version_info >= (3, 11): @@ -598,7 +604,6 @@ def __iter__(self): for i in iter_test(ast.For): assert i in (1, 2) - def test_decorator_cache_instruction(self): frame = inspect.currentframe()