From 0a77c96c0a4e069719d262cc3751f175e33c9d0d Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Sun, 29 Sep 2024 17:13:21 +0800
Subject: [PATCH] `next` can take default

---
 src/public/modules.c | 11 ++++++++---
 tests/51_yield.py    |  5 ++++-
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/src/public/modules.c b/src/public/modules.c
index d28d4432..9c021311 100644
--- a/src/public/modules.c
+++ b/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) {
diff --git a/tests/51_yield.py b/tests/51_yield.py
index 3e534880..880ea20e 100644
--- a/tests/51_yield.py
+++ b/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