diff options
Diffstat (limited to 'py/objset.c')
| -rw-r--r-- | py/objset.c | 19 |
1 files changed, 19 insertions, 0 deletions
diff --git a/py/objset.c b/py/objset.c index e41f2c47f..43bb61083 100644 --- a/py/objset.c +++ b/py/objset.c @@ -45,6 +45,24 @@ void set_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj print(env, "}"); } + +static mp_obj_t set_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { + mp_obj_set_t *o = lhs_in; + switch (op) { + case RT_COMPARE_OP_IN: + case RT_COMPARE_OP_NOT_IN: + { + mp_obj_t elem = mp_set_lookup(&o->set, rhs_in, false); + return ((op == RT_COMPARE_OP_IN) ^ (elem == NULL)) + ? mp_const_true : mp_const_false; + } + default: + // op not supported + return NULL; + } +} + + static mp_obj_t set_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args) { switch (n_args) { case 0: @@ -440,6 +458,7 @@ const mp_obj_type_t set_type = { { &mp_const_type }, "set", .print = set_print, + .binary_op = set_binary_op, .make_new = set_make_new, .binary_op = set_binary_op, .getiter = set_getiter, |