Quellcode durchsuchen

`next` can take default

blueloveTH vor 1 Jahr
Ursprung
Commit
0a77c96c0a
2 geänderte Dateien mit 12 neuen und 4 gelöschten Zeilen
  1. 8 3
      src/public/modules.c
  2. 4 1
      tests/51_yield.py

+ 8 - 3
src/public/modules.c

@@ -243,12 +243,17 @@ static bool builtins_iter(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);
     if(res == -1) return false;
     if(res) return true;
-    // StopIteration stored in py_retval()
-    return py_raise(py_retval());
+    if(argc == 1) {
+        // 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) {

+ 4 - 1
tests/51_yield.py

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