Skip to content

Commit be7f52c

Browse files
guilhermeleobasmarkc-614
authored andcommitted
Support next(iterator, default) (pytorch#159483)
Pull Request resolved: pytorch#159483 Approved by: https://github.com/mlazos ghstack dependencies: pytorch#159365, pytorch#159366, pytorch#159368
1 parent 5c70975 commit be7f52c

File tree

2 files changed

+7
-1
lines changed

2 files changed

+7
-1
lines changed

test/dynamo_expected_failures/CPython313-test_userlist-UserListTest.test_exhausted_iterator

Whitespace-only changes.

torch/_dynamo/variables/builtin.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
from ..exc import (
4747
AttributeMutationError,
4848
ObservedAttributeError,
49+
ObservedUserStopIteration,
4950
raise_observed_exception,
5051
unimplemented_v2,
5152
Unsupported,
@@ -2140,9 +2141,14 @@ def call_issubclass(self, tx: "InstructionTranslator", left_ty, right_ty):
21402141
def call_super(self, tx: "InstructionTranslator", a, b):
21412142
return variables.SuperVariable(a, b)
21422143

2143-
def call_next(self, tx: "InstructionTranslator", arg: VariableTracker):
2144+
def call_next(self, tx: "InstructionTranslator", *args):
2145+
arg = args[0]
21442146
try:
21452147
return arg.next_variable(tx)
2148+
except ObservedUserStopIteration:
2149+
if len(args) == 2:
2150+
return args[1]
2151+
raise
21462152
except Unsupported as ex:
21472153
if isinstance(arg, variables.BaseListVariable):
21482154
ex.remove_from_stats()

0 commit comments

Comments
 (0)