Fixed __contains__ comparison on lists, tuples and iter objects

Author: azmeuk

vm/src/obj/objiter.rs

@@ -7,6 +7,7 @@ use super::super::pyobject::{
     TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
+use super::objbool;
 use super::objstr;
 use super::objtype; // Required for arg_check! to use isinstance
 
@@ -61,13 +62,16 @@ fn iter_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     );
     loop {
         match vm.call_method(&iter, "__next__", vec![]) {
-            Ok(element) => {
-                if &element == needle {
-                    return Ok(vm.new_bool(true));
-                } else {
-                    continue;
+            Ok(element) => match vm.call_method(needle, "__eq__", vec![element.clone()]) {
+                Ok(value) => {
+                    if objbool::get_value(&value) {
+                        return Ok(vm.new_bool(true));
+                    } else {
+                        continue;
+                    }
                 }
-            }
+                Err(_) => return Err(vm.new_type_error("".to_string())),
+            },
             Err(_) => return Ok(vm.new_bool(false)),
         }
     }

vm/src/obj/objlist.rs

@@ -2,6 +2,7 @@ use super::super::pyobject::{
     AttributeProtocol, PyContext, PyFuncArgs, PyObjectKind, PyObjectRef, PyResult, TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
+use super::objbool;
 use super::objiter;
 use super::objsequence::{seq_equal, PySliceableSequence};
 use super::objstr;
@@ -162,16 +163,21 @@ fn reverse(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 }
 
-pub fn contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
-    trace!("list.len called with: {:?}", args);
+fn list_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    trace!("list.contains called with: {:?}", args);
     arg_check!(
         vm,
         args,
-        required = [(list, Some(vm.ctx.list_type())), (x, None)]
+        required = [(list, Some(vm.ctx.list_type())), (needle, None)]
     );
     for element in get_elements(list).iter() {
-        if x == element {
-            return Ok(vm.new_bool(true));
+        match vm.call_method(needle, "__eq__", vec![element.clone()]) {
+            Ok(value) => {
+                if objbool::get_value(&value) {
+                    return Ok(vm.new_bool(true));
+                }
+            }
+            Err(_) => return Err(vm.new_type_error("".to_string())),
         }
     }
 
@@ -181,7 +187,7 @@ pub fn contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 pub fn init(context: &PyContext) {
     let ref list_type = context.list_type;
     list_type.set_attr("__add__", context.new_rustfunc(list_add));
-    list_type.set_attr("__contains__", context.new_rustfunc(contains));
+    list_type.set_attr("__contains__", context.new_rustfunc(list_contains));
     list_type.set_attr("__eq__", context.new_rustfunc(list_eq));
     list_type.set_attr("__len__", context.new_rustfunc(list_len));
     list_type.set_attr("__new__", context.new_rustfunc(list_new));

vm/src/obj/objtuple.rs

@@ -2,6 +2,7 @@ use super::super::pyobject::{
     AttributeProtocol, PyContext, PyFuncArgs, PyObjectKind, PyObjectRef, PyResult, TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
+use super::objbool;
 use super::objsequence::seq_equal;
 use super::objstr;
 use super::objtype;
@@ -54,11 +55,16 @@ pub fn tuple_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
         args,
-        required = [(tuple, Some(vm.ctx.tuple_type())), (x, None)]
+        required = [(tuple, Some(vm.ctx.tuple_type())), (needle, None)]
     );
     for element in get_elements(tuple).iter() {
-        if x == element {
-            return Ok(vm.new_bool(true));
+        match vm.call_method(needle, "__eq__", vec![element.clone()]) {
+            Ok(value) => {
+                if objbool::get_value(&value) {
+                    return Ok(vm.new_bool(true));
+                }
+            }
+            Err(_) => return Err(vm.new_type_error("".to_string())),
         }
     }