diff options
Diffstat (limited to 'py')
| -rw-r--r-- | py/objset.c | 86 |
1 files changed, 76 insertions, 10 deletions
diff --git a/py/objset.c b/py/objset.c index 27cc91e4f8..2ed2abb611 100644 --- a/py/objset.c +++ b/py/objset.c @@ -9,6 +9,7 @@ #include "mpqstr.h" #include "obj.h" #include "runtime.h" +#include "runtime0.h" #include "map.h" typedef struct _mp_obj_set_t { @@ -227,7 +228,7 @@ static mp_obj_t set_isdisjoint(mp_obj_t self_in, mp_obj_t other) { } static MP_DEFINE_CONST_FUN_OBJ_2(set_isdisjoint_obj, set_isdisjoint); -static mp_obj_t set_issubset(mp_obj_t self_in, mp_obj_t other_in) { +static mp_obj_t set_issubset_internal(mp_obj_t self_in, mp_obj_t other_in, bool proper) { mp_obj_set_t *self; bool cleanup_self = false; if (MP_OBJ_IS_TYPE(self_in, &set_type)) { @@ -245,13 +246,17 @@ static mp_obj_t set_issubset(mp_obj_t self_in, mp_obj_t other_in) { other = set_make_new(NULL, 1, &other_in); cleanup_other = true; } - mp_obj_t iter = set_getiter(self); - mp_obj_t next; - mp_obj_t out = mp_const_true; - while ((next = set_it_iternext(iter)) != mp_const_stop_iteration) { - if (!mp_set_lookup(&other->set, next, MP_MAP_LOOKUP)) { - out = mp_const_false; - break; + bool out = true; + if (proper && self->set.used == other->set.used) { + out = false; + } else { + mp_obj_t iter = set_getiter(self); + mp_obj_t next; + while ((next = set_it_iternext(iter)) != mp_const_stop_iteration) { + if (!mp_set_lookup(&other->set, next, MP_MAP_LOOKUP)) { + out = false; + break; + } } } if (cleanup_self) { @@ -260,15 +265,39 @@ static mp_obj_t set_issubset(mp_obj_t self_in, mp_obj_t other_in) { if (cleanup_other) { set_clear(other); } - return out; + return MP_BOOL(out); +} +static mp_obj_t set_issubset(mp_obj_t self_in, mp_obj_t other_in) { + return set_issubset_internal(self_in, other_in, false); } static MP_DEFINE_CONST_FUN_OBJ_2(set_issubset_obj, set_issubset); +static mp_obj_t set_issubset_proper(mp_obj_t self_in, mp_obj_t other_in) { + return set_issubset_internal(self_in, other_in, true); +} + static mp_obj_t set_issuperset(mp_obj_t self_in, mp_obj_t other_in) { - return set_issubset(other_in, self_in); + return set_issubset_internal(other_in, self_in, false); } static MP_DEFINE_CONST_FUN_OBJ_2(set_issuperset_obj, set_issuperset); +static mp_obj_t set_issuperset_proper(mp_obj_t self_in, mp_obj_t other_in) { + return set_issubset_internal(other_in, self_in, true); +} + +static mp_obj_t set_equal(mp_obj_t self_in, mp_obj_t other_in) { + assert(MP_OBJ_IS_TYPE(self_in, &set_type)); + mp_obj_set_t *self = self_in; + if (!MP_OBJ_IS_TYPE(other_in, &set_type)) { + return mp_const_false; + } + mp_obj_set_t *other = other_in; + if (self->set.used != other->set.used) { + return mp_const_false; + } + return set_issubset(self_in, other_in); +} + static mp_obj_t set_pop(mp_obj_t self_in) { assert(MP_OBJ_IS_TYPE(self_in, &set_type)); mp_obj_set_t *self = self_in; @@ -341,6 +370,42 @@ 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_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) { + mp_obj_t args[] = {lhs, rhs}; + switch (op) { + case RT_BINARY_OP_OR: + return set_union(lhs, rhs); + case RT_BINARY_OP_XOR: + return set_symmetric_difference(lhs, rhs); + case RT_BINARY_OP_AND: + return set_intersect(lhs, rhs); + case RT_BINARY_OP_SUBTRACT: + return set_diff(2, args); + case RT_BINARY_OP_INPLACE_OR: + return set_union(lhs, rhs); + case RT_BINARY_OP_INPLACE_XOR: + return set_symmetric_difference(lhs, rhs); + case RT_BINARY_OP_INPLACE_AND: + return set_intersect(lhs, rhs); + case RT_BINARY_OP_INPLACE_SUBTRACT: + return set_diff(2, args); + case RT_COMPARE_OP_LESS: + return set_issubset_proper(lhs, rhs); + case RT_COMPARE_OP_MORE: + return set_issuperset_proper(lhs, rhs); + case RT_COMPARE_OP_EQUAL: + return set_equal(lhs, rhs); + case RT_COMPARE_OP_LESS_EQUAL: + return set_issubset(lhs, rhs); + case RT_COMPARE_OP_MORE_EQUAL: + return set_issuperset(lhs, rhs); + case RT_COMPARE_OP_NOT_EQUAL: + return MP_BOOL(set_equal(lhs, rhs) == mp_const_false); + default: + // op not supported + return NULL; + } +} /******************************************************************************/ /* set constructors & public C API */ @@ -372,6 +437,7 @@ const mp_obj_type_t set_type = { "set", .print = set_print, .make_new = set_make_new, + .binary_op = set_binary_op, .getiter = set_getiter, .methods = set_type_methods, }; |