py: Add optional support for 2-argument version of built-in next().

Configurable via MICROPY_PY_BUILTINS_NEXT2, disabled by default.

ports/unix/mpconfigport_coverage.h

@@ -39,6 +39,7 @@
 #define MICROPY_MODULE_GETATTR         (1)
 #define MICROPY_PY_DELATTR_SETATTR     (1)
 #define MICROPY_PY_REVERSE_SPECIAL_METHODS (1)
+#define MICROPY_PY_BUILTINS_NEXT2      (1)
 #define MICROPY_PY_BUILTINS_RANGE_BINOP (1)
 #define MICROPY_PY_BUILTINS_HELP       (1)
 #define MICROPY_PY_BUILTINS_HELP_MODULES (1)

py/builtin.h

@@ -63,7 +63,11 @@ MP_DECLARE_CONST_FUN_OBJ_1(mp_builtin_len_obj);
 MP_DECLARE_CONST_FUN_OBJ_0(mp_builtin_locals_obj);
 MP_DECLARE_CONST_FUN_OBJ_KW(mp_builtin_max_obj);
 MP_DECLARE_CONST_FUN_OBJ_KW(mp_builtin_min_obj);
+#if MICROPY_PY_BUILTINS_NEXT2
+MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_builtin_next_obj);
+#else
 MP_DECLARE_CONST_FUN_OBJ_1(mp_builtin_next_obj);
+#endif
 MP_DECLARE_CONST_FUN_OBJ_1(mp_builtin_oct_obj);
 MP_DECLARE_CONST_FUN_OBJ_1(mp_builtin_ord_obj);
 MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mp_builtin_pow_obj);

py/modbuiltins.c

@@ -316,6 +316,22 @@ MP_DEFINE_CONST_FUN_OBJ_KW(mp_builtin_min_obj, 1, mp_builtin_min);
 
 #endif
 
+#if MICROPY_PY_BUILTINS_NEXT2
+STATIC mp_obj_t mp_builtin_next(size_t n_args, const mp_obj_t *args) {
+    if (n_args == 1) {
+        mp_obj_t ret = mp_iternext_allow_raise(args[0]);
+        if (ret == MP_OBJ_STOP_ITERATION) {
+            nlr_raise(mp_obj_new_exception(&mp_type_StopIteration));
+        } else {
+            return ret;
+        }
+    } else {
+        mp_obj_t ret = mp_iternext(args[0]);
+        return ret == MP_OBJ_STOP_ITERATION ? args[1] : ret;
+    }
+}
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_builtin_next_obj, 1, 2, mp_builtin_next);
+#else
 STATIC mp_obj_t mp_builtin_next(mp_obj_t o) {
     mp_obj_t ret = mp_iternext_allow_raise(o);
     if (ret == MP_OBJ_STOP_ITERATION) {
@@ -325,6 +341,7 @@ STATIC mp_obj_t mp_builtin_next(mp_obj_t o) {
     }
 }
 MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_next_obj, mp_builtin_next);
+#endif
 
 STATIC mp_obj_t mp_builtin_oct(mp_obj_t o_in) {
     #if MICROPY_PY_BUILTINS_STR_OP_MODULO

py/mpconfig.h

@@ -881,6 +881,11 @@ typedef double mp_float_t;
 #define MICROPY_PY_BUILTINS_RANGE_BINOP (0)
 #endif
 
+// Support for callling next() with second argument
+#ifndef MICROPY_PY_BUILTINS_NEXT2
+#define MICROPY_PY_BUILTINS_NEXT2 (0)
+#endif
+
 // Whether to support rounding of integers (incl bignum); eg round(123,-1)=120
 #ifndef MICROPY_PY_BUILTINS_ROUND_INT
 #define MICROPY_PY_BUILTINS_ROUND_INT (0)

tests/basics/builtin_next_arg2.py

+# test next(iter, default)
+
+try:
+    next(iter([]), 42)
+except TypeError:  # 2-argument version not supported
+    print('SKIP')
+    raise SystemExit
+
+print(next(iter([]), 42))
+print(next(iter(range(0)), 42))
+print(next((x for x in [0] if x == 1), 43))
+
+def gen():
+    yield 1
+    yield 2
+
+g = gen()
+print(next(g, 42))
+print(next(g, 43))
+print(next(g, 44))
+
+class Gen:
+    def __init__(self):
+        self.b = False
+
+    def __next__(self):
+        if self.b:
+            raise StopIteration
+        self.b = True
+        return self.b
+
+g = Gen()
+print(next(g, 44))
+print(next(g, 45))