feat: support for __next__

alexmojaki · Sep 24, 2023 · 81888a3 · 81888a3
1 parent c668b55
commit 81888a3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git 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
@@ -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()