[pypy-commit] pypy numpy-ufuncs3: Add square

[taavi_burns]

Author: Taavi Burns <taavi.burns at gmail.com>
Branch: numpy-ufuncs3
Changeset: r54054:128dd4bc1819
Date: 2012-03-28 23:23 -0400
http://bitbucket.org/pypy/pypy/changeset/128dd4bc1819/

Log:	Add square

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
@@ -122,6 +122,7 @@
         ("sinh", "sinh"),
         ("subtract", "subtract"),
         ('sqrt', 'sqrt'),
+        ('square', 'square'),
         ("tan", "tan"),
         ("tanh", "tanh"),
         ('bitwise_and', 'bitwise_and'),
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
@@ -549,6 +549,7 @@
             ("expm1", "expm1", 1, {"promote_to_float": True}),
 
             ('sqrt', 'sqrt', 1, {'promote_to_float': True}),
+            ('square', 'square', 1, {'promote_to_float': True}),
 
             ("sin", "sin", 1, {"promote_to_float": True}),
             ("cos", "cos", 1, {"promote_to_float": True}),
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
@@ -455,6 +455,19 @@
         assert math.isnan(sqrt(-1))
         assert math.isnan(sqrt(nan))
 
+    def test_square(self):
+        import math
+        from _numpypy import square
+
+        nan, inf, ninf = float("nan"), float("inf"), float("-inf")
+
+        assert math.isnan(square(nan))
+        assert math.isinf(square(inf))
+        assert math.isinf(square(ninf))
+        assert square(ninf) > 0
+        assert [square(x) for x in range(-5, 5)] == [x*x for x in range(-5, 5)]
+        assert math.isinf(square(1e300))
+
     def test_radians(self):
         import math
         from _numpypy import radians, array
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
@@ -741,6 +741,10 @@
         except ValueError:
             return rfloat.NAN
 
+    @simple_unary_op
+    def square(self, v):
+        return v*v
+
     @raw_unary_op
     def isnan(self, v):
         return rfloat.isnan(v)