Merge pull request RustPython#479 from calixteman/fix_equality

Fix issue with equality

Author: windelbouwman

tests/snippets/list.py

@@ -74,3 +74,46 @@
 foo = bar = [1]
 foo += [2]
 assert (foo, bar) == ([1, 2], [1, 2])
+
+
+x = [1]
+x.append(x)
+assert x in x
+assert x.index(x) == 1
+assert x.count(x) == 1
+x.remove(x)
+assert x not in x
+
+class Foo(object):
+    def __eq__(self, x):
+        return False
+
+foo = Foo()
+foo1 = Foo()
+x = [1, foo, 2, foo, []]
+assert x == x
+assert foo in x
+assert 2 in x
+assert x.index(foo) == 1
+assert x.count(foo) == 2
+assert x.index(2) == 2
+assert [] in x
+assert x.index([]) == 4
+assert foo1 not in x
+x.remove(foo)
+assert x.index(foo) == 2
+assert x.count(foo) == 1
+
+x = []
+x.append(x)
+assert x == x
+
+a = [1, 2, 3]
+b = [1, 2, 3]
+c = [a, b]
+a.append(c)
+b.append(c)
+
+assert a == b
+
+assert [foo] == [foo]

tests/snippets/tuple.py

@@ -27,3 +27,10 @@
 
 assert (None, "", 1).index(1) == 2
 assert 1 in (None, "", 1)
+
+class Foo(object):
+    def __eq__(self, x):
+        return False
+
+foo = Foo()
+assert (foo,) == (foo,)

vm/src/obj/objlist.rs

@@ -1,5 +1,6 @@
 use super::super::pyobject::{
-    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
+    IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult,
+    TypeProtocol,
 };
 use super::super::vm::{ReprGuard, VirtualMachine};
 use super::objbool;
@@ -66,6 +67,10 @@ fn list_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(zelf, Some(vm.ctx.list_type())), (other, None)]
     );
 
+    if zelf.is(&other) {
+        return Ok(vm.ctx.new_bool(true));
+    }
+
     let result = if objtype::isinstance(other, &vm.ctx.list_type()) {
         let zelf = get_elements(zelf);
         let other = get_elements(other);
@@ -249,9 +254,13 @@ fn list_count(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let elements = get_elements(zelf);
     let mut count: usize = 0;
     for element in elements.iter() {
-        let is_eq = vm._eq(element.clone(), value.clone())?;
-        if objbool::boolval(vm, is_eq)? {
+        if value.is(&element) {
             count += 1;
+        } else {
+            let is_eq = vm._eq(element.clone(), value.clone())?;
+            if objbool::boolval(vm, is_eq)? {
+                count += 1;
+            }
         }
     }
     Ok(vm.context().new_int(count))
@@ -277,6 +286,9 @@ fn list_index(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(list, Some(vm.ctx.list_type())), (needle, None)]
     );
     for (index, element) in get_elements(list).iter().enumerate() {
+        if needle.is(&element) {
+            return Ok(vm.context().new_int(index));
+        }
         let py_equal = vm._eq(needle.clone(), element.clone())?;
         if objbool::get_value(&py_equal) {
             return Ok(vm.context().new_int(index));
@@ -350,6 +362,9 @@ fn list_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(list, Some(vm.ctx.list_type())), (needle, None)]
     );
     for element in get_elements(list).iter() {
+        if needle.is(&element) {
+            return Ok(vm.new_bool(true));
+        }
         match vm._eq(needle.clone(), element.clone()) {
             Ok(value) => {
                 if objbool::get_value(&value) {
@@ -431,9 +446,12 @@ fn list_remove(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(list, Some(vm.ctx.list_type())), (needle, None)]
     );
 
-    let mut elements = get_mut_elements(list);
     let mut ri: Option<usize> = None;
-    for (index, element) in elements.iter().enumerate() {
+    for (index, element) in get_elements(list).iter().enumerate() {
+        if needle.is(&element) {
+            ri = Some(index);
+            break;
+        }
         let py_equal = vm._eq(needle.clone(), element.clone())?;
         if objbool::get_value(&py_equal) {
             ri = Some(index);
@@ -442,6 +460,7 @@ fn list_remove(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 
     if let Some(index) = ri {
+        let mut elements = get_mut_elements(list);
         elements.remove(index);
         Ok(vm.get_none())
     } else {

vm/src/obj/objsequence.rs

@@ -1,4 +1,6 @@
-use super::super::pyobject::{PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol};
+use super::super::pyobject::{
+    IdProtocol, PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
+};
 use super::super::vm::VirtualMachine;
 use super::objbool;
 use super::objint;
@@ -178,10 +180,12 @@ pub fn seq_equal(
 ) -> Result<bool, PyObjectRef> {
     if zelf.len() == other.len() {
         for (a, b) in Iterator::zip(zelf.iter(), other.iter()) {
-            let eq = vm._eq(a.clone(), b.clone())?;
-            let value = objbool::boolval(vm, eq)?;
-            if !value {
-                return Ok(false);
+            if !a.is(&b) {
+                let eq = vm._eq(a.clone(), b.clone())?;
+                let value = objbool::boolval(vm, eq)?;
+                if !value {
+                    return Ok(false);
+                }
             }
         }
         Ok(true)