[pypy-commit] pypy numpy-comparison: Returning False for "arr == None". Preparations for "zeors(3) == zeros(4)" to

[snus_mumrik]

Author: Ilya Osadchiy <osadchiy.ilya at gmail.com>
Branch: numpy-comparison
Changeset: r48246:58762c30991d
Date: 2011-10-01 21:26 +0300
http://bitbucket.org/pypy/pypy/changeset/58762c30991d/

Log:	Returning False for "arr == None". Preparations for "zeors(3) ==
	zeros(4)" to return False

diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -74,8 +74,24 @@
     descr_pow = _binop_impl("power")
     descr_mod = _binop_impl("mod")
 
-    descr_eq = _binop_impl("equal")
-    descr_ne = _binop_impl("not_equal")
+    def _eq_ne_impl(ufunc_name, fallback_res):
+        assert isinstance(fallback_res, bool)
+        def impl(self, space, w_other):
+            # Unlike ufunc, operator may return simple bool
+            if space.is_w(w_other, space.w_None):
+                # Special case
+                return space.wrap(fallback_res)
+            try:
+                return getattr(interp_ufuncs.get(space), ufunc_name).call(space, [self, w_other])
+            except OperationError, e:
+                if e.match(space, space.w_ValueError):
+                    # For the case when arrays of incompatible size are compared
+                    return space.wrap(fallback_res)
+                raise
+        return func_with_new_name(impl, "binop_%s_impl" % ufunc_name)
+
+    descr_eq = _eq_ne_impl("equal", False)
+    descr_ne = _eq_ne_impl("not_equal", True)
     descr_lt = _binop_impl("less")
     descr_le = _binop_impl("less_equal")
     descr_gt = _binop_impl("greater")
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -577,6 +577,15 @@
             for i in xrange(5):
                 assert c[i] == func(b[i], 3)
 
+    def test_eq_ne(self):
+        from numpy import array
+        a = array(range(5))
+        assert (a == None) is False
+        assert (a != None) is True
+        # TODO: uncomment after size check is implemented
+        # b = array(range(2))
+        # assert (a == b) is False
+        # assert (a != b) is True
 
 class AppTestSupport(object):
     def setup_class(cls):