[pypy-commit] pypy default: Adds ones_like and zeros_like ufuncs.

[taavi_burns]

Author: Taavi Burns <taavi.burns at gmail.com>
Branch: 
Changeset: r66133:78d2ef3b1442
Date: 2013-08-13 15:51 -0400
http://bitbucket.org/pypy/pypy/changeset/78d2ef3b1442/

Log:	Adds ones_like and zeros_like ufuncs.

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -176,6 +176,8 @@
         ('logaddexp2', 'logaddexp2'),
         ('real', 'real'),
         ('imag', 'imag'),
+        ('ones_like', 'ones_like'),
+        ('zeros_like', 'zeros_like'),
     ]:
         interpleveldefs[exposed] = "interp_ufuncs.get(space).%s" % impl
 
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -676,6 +676,8 @@
                                        "allow_complex": False}),
             ("logaddexp2", "logaddexp2", 2, {"promote_to_float": True,
                                        "allow_complex": False}),
+            ("ones_like", "ones_like", 1),
+            ("zeros_like", "zeros_like", 1),
         ]:
             self.add_ufunc(space, *ufunc_def)
 
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -955,4 +955,14 @@
         assert logaddexp2(float('inf'), float('-inf')) == float('inf')
         assert logaddexp2(float('inf'), float('inf')) == float('inf')
 
+    def test_ones_like(self):
+        from numpypy import array, complex, ones_like
 
+        assert ones_like(2) == array(1)
+        assert ones_like(complex(2)) == array(complex(1))
+
+    def test_zeros_like(self):
+        from numpypy import array, complex, zeros_like
+
+        assert zeros_like(2) == array(0)
+        assert zeros_like(complex(2)) == array(complex(0))
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -314,6 +314,15 @@
         else:
             return v
 
+    @simple_unary_op
+    def ones_like(self, v):
+        return 1
+
+    @simple_unary_op
+    def zeros_like(self, v):
+        return 0
+
+
 class NonNativePrimitive(Primitive):
     _mixin_ = True
 
@@ -1609,6 +1618,15 @@
         except ValueError:
             return rfloat.NAN, rfloat.NAN
 
+    @complex_unary_op
+    def ones_like(self, v):
+        return 1, 0
+
+    @complex_unary_op
+    def zeros_like(self, v):
+        return 0, 0
+
+
 class Complex64(ComplexFloating, BaseType):
     _attrs_ = ()