Commits

Maciej Fijalkowski committed 103077d Draft Merge

Merged in redorlik/pypy/numpypy_count_nonzero (pull request #76)

  • Participants
  • Parent commits ad478c8, 1be7eb1

Comments (0)

Files changed (7)

File pypy/module/micronumpy/__init__.py

         'eye': 'app_numpy.eye',
         'max': 'app_numpy.max',
         'arange': 'app_numpy.arange',
+        'count_nonzero': 'app_numpy.count_nonzero',
     }

File pypy/module/micronumpy/app_numpy.py

 
 import _numpypy
 
+def count_nonzero(a):
+    if not hasattr(a, 'count_nonzero'):
+        a = _numpypy.array(a)
+    return a.count_nonzero()
 
 def average(a):
     # This implements a weighted average, for now we don't implement the

File pypy/module/micronumpy/compile.py

     pass
 
 SINGLE_ARG_FUNCTIONS = ["sum", "prod", "max", "min", "all", "any",
-                        "unegative", "flat", "tostring"]
+                        "unegative", "flat", "tostring","count_nonzero"]
 TWO_ARG_FUNCTIONS = ["dot", 'take']
 THREE_ARG_FUNCTIONS = ['where']
 
             elif self.name == "tostring":
                 arr.descr_tostring(interp.space)
                 w_res = None
+            elif self.name == "count_nonzero":
+                w_res = arr.descr_count_nonzero(interp.space)
             else:
                 assert False # unreachable code
         elif self.name in TWO_ARG_FUNCTIONS:
             return w_res
         if isinstance(w_res, FloatObject):
             dtype = get_dtype_cache(interp.space).w_float64dtype
+        elif isinstance(w_res, IntObject):
+            dtype = get_dtype_cache(interp.space).w_int64dtype
         elif isinstance(w_res, BoolObject):
             dtype = get_dtype_cache(interp.space).w_booldtype
         elif isinstance(w_res, interp_boxes.W_GenericBox):

File pypy/module/micronumpy/interp_numarray.py

                 i += 1
         return Chunks(result)
 
+    def descr_count_nonzero(self, space):
+        concr = self.get_concrete()
+        res = concr.count_all_true()
+        return space.wrap(res)
+
     def count_all_true(self):
         sig = self.find_sig()
         frame = sig.create_frame(self)
     take = interp2app(BaseArray.descr_take),
     compress = interp2app(BaseArray.descr_compress),
     repeat = interp2app(BaseArray.descr_repeat),
+    count_nonzero = interp2app(BaseArray.descr_count_nonzero),
 )
 
 

File pypy/module/micronumpy/test/test_numarray.py

         raises(ValueError, "array(5).item(1)")
         assert array([1]).item() == 1
 
+    def test_count_nonzero(self):
+        from _numpypy import array
+        a = array([1,0,5,0,10])
+        assert a.count_nonzero() == 3
+ 
+
 class AppTestSupport(BaseNumpyAppTest):
     def setup_class(cls):
         import struct

File pypy/module/micronumpy/test/test_ufuncs.py

         raises(ValueError, count_reduce_items, a, -4)
         raises(ValueError, count_reduce_items, a, (0, 2, -4))
 
+    def test_count_nonzero(self):
+        from _numpypy import where, count_nonzero, arange
+        a = arange(10)
+        assert count_nonzero(a) == 9
+        a[9] = 0
+        assert count_nonzero(a) == 8
+
     def test_true_divide(self):
         from _numpypy import arange, array, true_divide
         assert (true_divide(arange(3), array([2, 2, 2])) == array([0, 0.5, 1])).all()

File pypy/module/micronumpy/test/test_zjit.py

                                 'int_sub': 3,
                                 'jump': 1,
                                 'setinteriorfield_raw': 1})
+    
+    def define_count_nonzero():
+        return """
+        a = [[0, 2, 3, 4], [5, 6, 0, 8], [9, 10, 11, 0]]
+        count_nonzero(a) 
+        """
+
+    def test_count_nonzero(self):
+        result = self.run("count_nonzero")
+        assert result == 9
+        self.check_simple_loop({'setfield_gc': 3, 
+                                'getinteriorfield_raw': 1, 
+                                'guard_false': 1, 
+                                'jump': 1, 
+                                'int_ge': 1, 
+                                'new_with_vtable': 1, 
+                                'int_add': 2, 
+                                'float_ne': 1})
+