Implement simple 'a in b' support

src/analysis/type_analysis.cpp

@@ -284,10 +284,12 @@ class BasicBlockTypePropagator : public ExprVisitor, public StmtVisitor {
             CompilerType *left = getType(node->left);
             CompilerType *right = getType(node->comparators[0]);
 
-            if (node->ops[0] == AST_TYPE::Is || node->ops[0] == AST_TYPE::IsNot) {
+            AST_TYPE::AST_TYPE op_type = node->ops[0];
+            if (op_type == AST_TYPE::Is || op_type == AST_TYPE::IsNot || op_type == AST_TYPE::In || op_type == AST_TYPE::NotIn) {
                 assert(node->ops.size() == 1 && "I don't think this should happen");
                 return BOOL;
             }
+
             std::string name = getOpName(node->ops[0]);
             CompilerType *attr_type = left->getattrType(&name, true);
 

src/core/ast.cpp

@@ -67,6 +67,8 @@ std::string getOpSymbol(int op_type) {
             return "not";
         case AST_TYPE::NotEq:
             return "!=";
+        case AST_TYPE::NotIn:
+            return "not in";
         case AST_TYPE::Pow:
             return "**";
         case AST_TYPE::RShift:

src/runtime/list.cpp

@@ -16,6 +16,7 @@
 #include <sstream>
 #include <algorithm>
 
+#include "core/ast.h"
 #include "core/common.h"
 #include "core/stats.h"
 #include "core/types.h"
@@ -276,6 +277,20 @@ Box* listSort1(BoxedList* self) {
     return None;
 }
 
+Box* listContains(BoxedList* self, Box *elt) {
+    int size = self->size;
+    for (int i = 0; i < size; i++) {
+        Box* e = self->elts->elts[i];
+        Box* cmp = compareInternal(e, elt, AST_TYPE::Eq, NULL);
+        bool b = nonzero(cmp);
+        if (b)
+            return True;
+    }
+    return False;
+}
+
+
+
 BoxedClass *list_iterator_cls = NULL;
 extern "C" void listIteratorGCHandler(GCVisitor *v, void* p) {
     boxGCHandler(v, p);
@@ -347,6 +362,7 @@ void setupList() {
     list_cls->giveAttr("__add__", new BoxedFunction(boxRTFunction((void*)listAdd, NULL, 2, false)));
 
     list_cls->giveAttr("sort", new BoxedFunction(boxRTFunction((void*)listSort1, NULL, 1, false)));
+    list_cls->giveAttr("__contains__", new BoxedFunction(boxRTFunction((void*)listContains, BOXED_BOOL, 2, false)));
 
     CLFunction *new_ = boxRTFunction((void*)listNew1, NULL, 1, false);
     addRTFunction(new_, (void*)listNew2, NULL, 2, false);

src/runtime/objmodel.cpp

@@ -1758,6 +1758,30 @@ Box* compareInternal(Box* lhs, Box* rhs, int op_type, CompareRewriteArgs *rewrit
         return boxBool((lhs == rhs) ^ neg);
     }
 
+    if (op_type == AST_TYPE::In || op_type == AST_TYPE::NotIn) {
+        // TODO do rewrite
+
+        static const std::string str_contains("__contains__");
+        Box* contained = callattrInternal1(rhs, &str_contains, CLASS_ONLY, NULL, 1, lhs);
+        if (contained == NULL) {
+            static const std::string str_iter("__iter__");
+            Box* iter = callattrInternal0(rhs, &str_iter, CLASS_ONLY, NULL, 0);
+            if (iter)
+                ASSERT(isUserDefined(rhs->cls), "%s should probably have a __contains__", getTypeName(rhs)->c_str());
+            RELEASE_ASSERT(iter == NULL, "need to try iterating");
+
+            Box* getitem = getattr_internal(rhs, "__getitem__", false, false, NULL, NULL);
+            if (getitem)
+                ASSERT(isUserDefined(rhs->cls), "%s should probably have a __contains__", getTypeName(rhs)->c_str());
+            RELEASE_ASSERT(getitem == NULL, "need to try old iteration protocol");
+        }
+
+        bool b = nonzero(contained);
+        if (op_type == AST_TYPE::NotIn)
+            return boxBool(!b);
+        return boxBool(b);
+    }
+
     // Can do the guard checks after the Is/IsNot handling, since that is
     // irrespective of the object classes
     if (rewrite_args) {

test/tests/comparisons_failing.py → test/tests/comparisons_chained.py

@@ -6,7 +6,5 @@ def f(n):
     print "f(%d)" % n
     return n
 
-f(1) <= f(2) < f(3)
-
-for i in xrange(1, 4):
-    print i in range(6), i not in range(5)
+# f(3) shouldn't get called:
+f(1) <= f(2) < f(1) < f(3)

test/tests/contains.py

+class C(object):
+    def __init__(self, x):
+        self.x = x
+
+    def __eq__(self, y):
+        print "__eq__", y
+        return self.x
+
+    def __contains__(self, y):
+        print "__contains__", y
+        return self.x
+
+print 1 in C("hello") # "a in b" expressions get coerced to boolean
+print 2 in C("")
+
+print 1 in [C("hello")] # True
+print 2 in [C("")] # False
+
+for i in xrange(1, 4):
+    print i in range(6), i not in range(5)

test/tests/contains_getitem_fallback.py

+# expected: fail
+# - exceptions
+
+class D(object):
+    def __getitem__(self, idx):
+        print "getitem", idx
+
+        if idx >= 20:
+            raise IndexError()
+        return idx
+
+print 10 in D()
+print 1000 in D()
+
+