diff options
| -rw-r--r-- | ports/unix/mpconfigport_coverage.h | 1 | ||||
| -rw-r--r-- | py/builtin.h | 4 | ||||
| -rw-r--r-- | py/modbuiltins.c | 17 | ||||
| -rw-r--r-- | py/mpconfig.h | 5 | ||||
| -rw-r--r-- | tests/basics/builtin_next_arg2.py | 34 |
5 files changed, 61 insertions, 0 deletions
diff --git a/ports/unix/mpconfigport_coverage.h b/ports/unix/mpconfigport_coverage.h index 9ab442ff9..2519482b0 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 6f8964a25..2066c0617 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 c4de325c1..f0d0421d6 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 68a855f19..48427c3e5 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 000000000..b00155ec5 --- /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)) |