ENH: Object loop for bit_count

ganesh-k13 · ganesh-k13 · commit ab72dfc0afd7 · 2021-07-03T08:28:35.000+05:30
[2/2] `popcount` object loop changes
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -3537,7 +3537,7 @@ arcsinh: _UFunc_Nin1_Nout1[L['arcsinh'], L[8], None]
 arctan2: _UFunc_Nin2_Nout1[L['arctan2'], L[5], None]
 arctan: _UFunc_Nin1_Nout1[L['arctan'], L[8], None]
 arctanh: _UFunc_Nin1_Nout1[L['arctanh'], L[8], None]
-bit_count: _UFunc_Nin1_Nout1[L['bit_count'], L[10], None]
+bit_count: _UFunc_Nin1_Nout1[L['bit_count'], L[11], None]
 bitwise_and: _UFunc_Nin2_Nout1[L['bitwise_and'], L[12], L[-1]]
 bitwise_not: _UFunc_Nin1_Nout1[L['invert'], L[12], None]
 bitwise_or: _UFunc_Nin2_Nout1[L['bitwise_or'], L[12], L[0]]
diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
@@ -960,6 +960,7 @@ def english_upper(s):
           docstrings.get('numpy.core.umath.bit_count'),
           None,
           TD(ints),
+          TD('O', f='npy_ObjectPopCount'),
           ),
 'matmul' :
     Ufunc(2, 1, None,
diff --git a/numpy/core/src/umath/funcs.inc.src b/numpy/core/src/umath/funcs.inc.src
@@ -267,6 +267,38 @@ npy_ObjectClip(PyObject *arr, PyObject *min, PyObject *max) {
     return o;
 }
 
+static PyObject *
+npy_ObjectPopCount(PyObject *obj) {
+    PyObject *result;
+
+    /* Try to use inbuilt popcount if available */
+    {
+        static PyObject *builtin_popcount_func = NULL;
+
+        builtin_popcount_func = PyObject_GetAttrString(obj, "bit_count");
+        if (builtin_popcount_func != NULL) {
+            result = PyObject_CallFunction(builtin_popcount_func, NULL);
+            return result;
+        }
+        /* silence errors, and fall back on pure-python `bit_count` */
+        PyErr_Clear();
+    }
+
+    /* Otherwise, use our internal one, written in python */
+    {
+        static PyObject *internal_popcount_func = NULL;
+
+        npy_cache_import("numpy.core._internal", "_bit_count", &internal_popcount_func);
+        if (internal_popcount_func == NULL) {
+            return NULL;
+        }
+
+        result = PyObject_CallFunction(internal_popcount_func, "O", obj);
+
+        return result;
+    }
+}
+
 /*
  *****************************************************************************
  **                           COMPLEX FUNCTIONS                             **
