diff options
Diffstat (limited to 'py/objset.c')
-rw-r--r-- | py/objset.c | 19 |
1 files changed, 13 insertions, 6 deletions
diff --git a/py/objset.c b/py/objset.c index 439c6e96e6..222f76e405 100644 --- a/py/objset.c +++ b/py/objset.c @@ -32,7 +32,7 @@ STATIC void set_print(void (*print)(void *env, const char *fmt, ...), void *env, bool first = true; print(env, "{"); for (int i = 0; i < self->set.alloc; i++) { - if (self->set.table[i] != MP_OBJ_NULL) { + if (self->set.table[i] != MP_OBJ_NULL && self->set.table[i] != MP_OBJ_SENTINEL) { if (!first) { print(env, ", "); } @@ -83,7 +83,7 @@ STATIC mp_obj_t set_it_iternext(mp_obj_t self_in) { mp_obj_t *table = self->set->set.table; for (machine_uint_t i = self->cur; i < max; i++) { - if (table[i] != NULL) { + if (table[i] != MP_OBJ_NULL && table[i] != MP_OBJ_SENTINEL) { self->cur = i + 1; return table[i]; } @@ -307,12 +307,10 @@ STATIC mp_obj_t set_equal(mp_obj_t self_in, mp_obj_t other_in) { STATIC mp_obj_t set_pop(mp_obj_t self_in) { assert(MP_OBJ_IS_TYPE(self_in, &mp_type_set)); mp_obj_set_t *self = self_in; - - if (self->set.used == 0) { + mp_obj_t obj = mp_set_remove_first(&self->set); + if (obj == MP_OBJ_NULL) { nlr_jump(mp_obj_new_exception_msg(&mp_type_KeyError, "pop from an empty set")); } - mp_obj_t obj = mp_set_lookup(&self->set, NULL, - MP_MAP_LOOKUP_REMOVE_IF_FOUND | MP_MAP_LOOKUP_FIRST); return obj; } STATIC MP_DEFINE_CONST_FUN_OBJ_1(set_pop_obj, set_pop); @@ -375,6 +373,14 @@ STATIC mp_obj_t set_union(mp_obj_t self_in, mp_obj_t other_in) { } STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_union_obj, set_union); +STATIC mp_obj_t set_unary_op(int op, mp_obj_t self_in) { + mp_obj_set_t *self = self_in; + switch (op) { + case MP_UNARY_OP_BOOL: return MP_BOOL(self->set.used != 0); + case MP_UNARY_OP_LEN: return MP_OBJ_NEW_SMALL_INT((machine_int_t)self->set.used); + default: return MP_OBJ_NULL; // op not supported for None + } +} STATIC mp_obj_t set_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) { mp_obj_t args[] = {lhs, rhs}; @@ -450,6 +456,7 @@ const mp_obj_type_t mp_type_set = { .name = MP_QSTR_set, .print = set_print, .make_new = set_make_new, + .unary_op = set_unary_op, .binary_op = set_binary_op, .getiter = set_getiter, .locals_dict = (mp_obj_t)&set_locals_dict, |