From 42863830be2d19c7dcdf7ccf1fa66168b1bdc13a Mon Sep 17 00:00:00 2001
From: stijn <stijn@ignitron.net>
Date: Thu, 3 Jan 2019 15:19:42 +0100
Subject: [PATCH] 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 |  1 +
 py/builtin.h                       |  4 ++++
 py/modbuiltins.c                   | 17 +++++++++++++++
 py/mpconfig.h                      |  5 +++++
 tests/basics/builtin_next_arg2.py  | 34 ++++++++++++++++++++++++++++++
 5 files changed, 61 insertions(+)
 create mode 100644 tests/basics/builtin_next_arg2.py

diff --git a/ports/unix/mpconfigport_coverage.h b/ports/unix/mpconfigport_coverage.h
index 9ab442ff92..2519482b01 100644
--- a/ports/unix/mpconfigport_coverage.h
+++ b/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)
diff --git a/py/builtin.h b/py/builtin.h
index 6f8964a250..2066c06173 100644
--- a/py/builtin.h
+++ b/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);
diff --git a/py/modbuiltins.c b/py/modbuiltins.c
index c4de325c14..f0d0421d6c 100644
--- a/py/modbuiltins.c
+++ b/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
diff --git a/py/mpconfig.h b/py/mpconfig.h
index 68a855f199..48427c3e58 100644
--- a/py/mpconfig.h
+++ b/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)
diff --git a/tests/basics/builtin_next_arg2.py b/tests/basics/builtin_next_arg2.py
new file mode 100644
index 0000000000..b00155ec54
--- /dev/null
+++ b/tests/basics/builtin_next_arg2.py
@@ -0,0 +1,34 @@
+# 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))