BUG: Any dtype should call square on arr ** 2 (#29392)

MaanasArora · web-flow · commit 8f68377c0e90 · 2025-07-22T12:31:01.000+02:00
* BUG: update fast_scalar_power to handle special-case squaring for any array type except object arrays

* BUG: fix missing declaration

* TST: add test to ensure `arr**2` calls square for structured dtypes

* STY: remove whitespace

* BUG: replace new variable `is_square` with direct op comparison in `fast_scalar_power` function
diff --git a/numpy/_core/src/multiarray/number.c b/numpy/_core/src/multiarray/number.c
@@ -332,6 +332,7 @@ static int
 fast_scalar_power(PyObject *o1, PyObject *o2, int inplace, PyObject **result)
 {
     PyObject *fastop = NULL;
+
     if (PyLong_CheckExact(o2)) {
         int overflow = 0;
         long exp = PyLong_AsLongAndOverflow(o2, &overflow);
@@ -363,7 +364,12 @@ fast_scalar_power(PyObject *o1, PyObject *o2, int inplace, PyObject **result)
     }
 
     PyArrayObject *a1 = (PyArrayObject *)o1;
-    if (!(PyArray_ISFLOAT(a1) || PyArray_ISCOMPLEX(a1))) {
+    if (PyArray_ISOBJECT(a1)) {
+        return 1;
+    }
+    if (fastop != n_ops.square && !PyArray_ISFLOAT(a1) && !PyArray_ISCOMPLEX(a1)) {
+        // we special-case squaring for any array type
+        // gh-29388
         return 1;
     }
 
diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
@@ -4215,6 +4215,13 @@ def pow_for(exp, arr):
         assert_equal(obj_arr ** -1, pow_for(-1, obj_arr))
         assert_equal(obj_arr ** 2, pow_for(2, obj_arr))
 
+    def test_pow_calls_square_structured_dtype(self):
+        # gh-29388
+        dt = np.dtype([('a', 'i4'), ('b', 'i4')])
+        a = np.array([(1, 2), (3, 4)], dtype=dt)
+        with pytest.raises(TypeError, match="ufunc 'square' not supported"):
+            a ** 2
+
     def test_pos_array_ufunc_override(self):
         class A(np.ndarray):
             def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):

Original file line number	Diff line number	Diff line change
`@@ -332,6 +332,7 @@ static int`
`332`	`332`	`fast_scalar_power(PyObject o1, PyObject o2, int inplace, PyObject **result)`
`333`	`333`	`{`
`334`	`334`	`PyObject *fastop = NULL;`
	`335`	`+`
`335`	`336`	`if (PyLong_CheckExact(o2)) {`
`336`	`337`	`int overflow = 0;`
`337`	`338`	`long exp = PyLong_AsLongAndOverflow(o2, &overflow);`
`@@ -363,7 +364,12 @@ fast_scalar_power(PyObject o1, PyObject o2, int inplace, PyObject **result)`
`363`	`364`	`}`
`364`	`365`
`365`	`366`	`PyArrayObject a1 = (PyArrayObject )o1;`
`366`		`- if (!(PyArray_ISFLOAT(a1) \|\| PyArray_ISCOMPLEX(a1))) {`
	`367`	`+ if (PyArray_ISOBJECT(a1)) {`
	`368`	`+ return 1;`
	`369`	`+ }`
	`370`	`+ if (fastop != n_ops.square && !PyArray_ISFLOAT(a1) && !PyArray_ISCOMPLEX(a1)) {`
	`371`	`+ // we special-case squaring for any array type`
	`372`	`+ // gh-29388`
`367`	`373`	`return 1;`
`368`	`374`	`}`
`369`	`375`