added np.size() overload and added tests

sklam · Apr 16, 2024 · Mar 19, 2024 · Mar 19, 2024 · Mar 26, 2024 · Mar 26, 2024
commit 6aea86fcac126ff4a8285a88ec08150c316e4433
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
@@ -2555,6 +2555,15 @@ def impl(a):
         return np.asarray(a).shape
     return impl
 
+@overload(np.size)
+def np_size(a):
+    if not type_can_asarray(a):
+        raise errors.TypingError("The argument to np.size must be array-like")
+
+    def impl(a):
+        return np.asarray(a).size
+    return impl
+
 # ------------------------------------------------------------------------------
 
 

diff --git a/numba/tests/test_array_manipulation.py b/numba/tests/test_array_manipulation.py
@@ -174,6 +174,9 @@ def numpy_fill_diagonal(arr, val, wrap=False):
 def numpy_shape(arr):
     return np.shape(arr)
 
+def numpy_size(arr):
+    return np.size(arr)
+
 
 def numpy_flatnonzero(a):
     return np.flatnonzero(a)

diff --git a/numba/tests/test_record_dtype.py b/numba/tests/test_record_dtype.py
@@ -332,6 +332,10 @@ def get_shape(rec):
     return np.shape(rec.j)
 
 
+def get_size(rec):
+    return np.size(rec.j)
+
+
 def get_charseq(ary, i):
     return ary[i].n