Implemented support for `in` and `not in` operators.

1 files changed, 20 insertions, 0 deletions

diff --git a/py/objset.c b/py/objset.c
index 67dab11df..71ed55335 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 @@ 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:
@@ -58,6 +77,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,
 };