[pypy-commit] pypy numpypy-frompyfunc: more tests, and a failing one

[mattip]

Author: mattip
Branch: numpypy-frompyfunc
Changeset: r50518:4136f83fdb42
Date: 2011-12-13 21:52 +0200
http://bitbucket.org/pypy/pypy/changeset/4136f83fdb42/

Log:	more tests, and a failing one

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
@@ -26,6 +26,10 @@
     greens=['shapelen', 'signature'],
     reds=['self', 'source', 'source_iter', 'res_iter']
 )
+ufunc_driver = jit.JitDriver(
+    greens=['shapelen', 'signature'],
+    reds=['result_size', 'i', 'ri', 'arr', 'result']
+)
 
 def _find_shape_and_elems(space, w_iterable):
     shape = [space.len_w(w_iterable)]
@@ -1571,10 +1575,10 @@
         shapelen = len(arr.shape)
         result_size = arr.find_size()
         while not ri.done():
-            #numpy_driver.jit_merge_point(signature=signature,
-            #                             shapelen=shapelen,
-            #                             result_size=result_size, i=i, ri=ri,
-            #                             self=self, result=result)
+            ufunc_driver.jit_merge_point(signature=signature,
+                                         shapelen=shapelen,
+                                         result_size=result_size, i=i, ri=ri,
+                                         arr=arr, result=result)
             result.dtype.setitem(result.storage, ri.offset,
                                 space.call_function(self.w_func, arr.eval(i)))
             i = i.next(shapelen)
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
@@ -380,8 +380,31 @@
         assert (ufunc([-1, 0, 3, 15]) == [1, 0, 3, 15]).all()
 
     def test_frompyfunc_foo(self):
+        from numpypy import frompyfunc, array
         def foo(x):
             return x * x + 1
+        def bar(x):
+            return x + 1
+        ufunc = frompyfunc(foo, 1, 1)
+        assert (ufunc(range(10)) == array(range(10)) * range(10) + 1).all()
+        #Make sure the user-visible function does not modify the ufunc
+        foo = bar
+        assert (ufunc(range(10)) == array(range(10)) * range(10) + 1).all()
+        #but messing with the func_code WILL change it: numpy is sensitive
+        #to this in the same way
+        def foo(x):
+            return x * x + 1
+        def bar(x):
+            return x + 1
         from numpypy import frompyfunc, array
         ufunc = frompyfunc(foo, 1, 1)
         assert (ufunc(range(10)) == array(range(10)) * range(10) + 1).all()
+        foo.func_code = bar.func_code
+        assert not (ufunc(range(10)) == array(range(10)) * range(10) + 1).all()
+    def test_frompyfunc_broadcast(self):
+        from numpypy import frompyfunc, array
+        def foo(x, y):
+            return x * y + 1
+        ufunc = frompyfunc(foo, 2, 1)
+        assert (ufunc(range(10),range(10)) == array(range(10)) * range(10) + 1).all()
+