Implemented support for `in` and `not in` operators.
diff --git a/py/objset.c b/py/objset.c
index 67dab11..71ed553 100644
--- a/py/objset.c
+++ b/py/objset.c
@@ -8,6 +8,7 @@
#include "mpqstr.h"
#include "obj.h"
#include "runtime.h"
+#include "runtime0.h"
#include "map.h"
typedef struct _mp_obj_set_t {
@@ -31,6 +32,24 @@
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:
@@ -58,6 +77,7 @@
{ &mp_const_type },
"set",
.print = set_print,
+ .binary_op = set_binary_op,
.make_new = set_make_new,
};