Author: mattip <matti.pi...@gmail.com>
Branch: 
Changeset: r80221:de3e116ddd1c
Date: 2015-10-14 23:40 +0300
http://bitbucket.org/pypy/pypy/changeset/de3e116ddd1c/

Log:    test and fix (via a fastpath for scalars) SliceArrays with scalar
        assignment

diff --git a/pypy/module/micronumpy/concrete.py 
b/pypy/module/micronumpy/concrete.py
--- a/pypy/module/micronumpy/concrete.py
+++ b/pypy/module/micronumpy/concrete.py
@@ -70,7 +70,10 @@
 
     @jit.unroll_safe
     def setslice(self, space, arr):
-        if len(arr.get_shape()) >  len(self.get_shape()):
+        if arr.get_size() == 1:
+            # we can always set self[:] = scalar
+            pass
+        elif len(arr.get_shape()) >  len(self.get_shape()):
             # record arrays get one extra dimension
             if not self.dtype.is_record() or \
                     len(arr.get_shape()) > len(self.get_shape()) + 1:
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
@@ -319,6 +319,28 @@
             assert out0.dtype in (int, complex) 
             assert (out0 == in0 * 2).all()
 
+    def test_frompyfunc_scalar(self):
+        import sys
+        import numpy as np
+        if '__pypy__' not in sys.builtin_module_names:
+            skip('PyPy only frompyfunc extension')
+
+        def summer(in0):
+            out = np.empty(1, in0.dtype)
+            out[0] = in0.sum()
+            return out
+
+        pysummer = np.frompyfunc([summer, summer], 1, 1,
+                            dtypes=[np.dtype(int), np.dtype(int),
+                                np.dtype(complex), np.dtype(complex)],
+                            stack_inputs=False, signature='(m,m)->()',
+                          )
+        for d in [np.dtype(float), np.dtype('uint8'), np.dtype('complex64')]:
+            in0 = np.arange(4, dtype=d).reshape(1, 2, 2)
+            out0 = pysummer(in0)
+            assert out0 == in0.sum()
+            assert out0.dtype in (int, complex)
+
     def test_ufunc_kwargs(self):
         from numpy import ufunc, frompyfunc, arange, dtype
         def adder(a, b):
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to