next can take default

This commit is contained in:
blueloveTH 2024-09-29 17:13:21 +08:00
parent 571a080127
commit 0a77c96c0a
2 changed files with 12 additions and 4 deletions

View File

@ -243,12 +243,17 @@ static bool builtins_iter(int argc, py_Ref argv) {
} }
static bool builtins_next(int argc, py_Ref argv) { static bool builtins_next(int argc, py_Ref argv) {
PY_CHECK_ARGC(1); if(argc == 0 || argc > 2) return TypeError("next() takes 1 or 2 arguments");
int res = py_next(argv); int res = py_next(argv);
if(res == -1) return false; if(res == -1) return false;
if(res) return true; if(res) return true;
// StopIteration stored in py_retval() if(argc == 1) {
return py_raise(py_retval()); // StopIteration stored in py_retval()
return py_raise(py_retval());
} else {
py_assign(py_retval(), py_arg(1));
return true;
}
} }
static bool builtins_hash(int argc, py_Ref argv) { static bool builtins_hash(int argc, py_Ref argv) {

View File

@ -4,7 +4,7 @@ def g():
a = g() a = g()
assert next(a) == 1 assert next(a) == 1
assert next(a) == 2 assert next(a, None) == 2
try: try:
next(a) next(a)
@ -12,6 +12,9 @@ try:
except StopIteration: except StopIteration:
pass pass
assert next(a, 3) == 3
assert next(a, 4) == 4
def f(n): def f(n):
for i in range(n): for i in range(n):
yield i yield i