Merge pull request #142 from chipaca/containment

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

1 files changed, 8 insertions, 0 deletions

diff --git a/py/objset.c b/py/objset.c
index e41f2c47f..ba083ada8 100644
--- a/py/objset.c
+++ b/py/objset.c
@@ -45,6 +45,7 @@ void set_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj
     print(env, "}");
 }
 
+
 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:
@@ -405,6 +406,13 @@ static mp_obj_t set_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
         return set_issuperset(lhs, rhs);
     case RT_COMPARE_OP_NOT_EQUAL:
         return MP_BOOL(set_equal(lhs, rhs) == mp_const_false);
+    case RT_COMPARE_OP_IN:
+    case RT_COMPARE_OP_NOT_IN:
+    {
+        mp_obj_set_t *o = lhs;
+        mp_obj_t elem = mp_set_lookup(&o->set, rhs, MP_MAP_LOOKUP);
+        return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (elem == NULL));
+    }
     default:
         // op not supported
         return NULL;