gh-117999: fixed small nonnegative integer powers of complex numbers

Lib/test/test_complex.py

@@ -5,8 +5,11 @@
                                INVALID_UNDERSCORE_LITERALS)
 
 from random import random
-from math import isnan, copysign
+from math import atan2, isnan, copysign
+from functools import reduce
+from itertools import combinations
 import operator
+import _testcapi
 
 INF = float("inf")
 NAN = float("nan")
@@ -351,6 +354,45 @@ def test_pow_with_small_integer_exponents(self):
                     self.assertEqual(str(float_pow), str(int_pow))
                     self.assertEqual(str(complex_pow), str(int_pow))
 
+        # Check that complex numbers with special components
+        # are correctly handled.
+        values = [complex(*_) for _ in combinations([1, -1, 0.0, -0.0, 2,
+                                                     -3, INF, -INF, NAN], 2)]
+        exponents = [0, 1, 2, 3, 4, 5, 6, 19]
+        for z in values:
+            for e in exponents:
+                try:
+                    r_pow = z**e
+                except OverflowError:
+                    continue
+                r_pro = reduce(lambda x, y: x*y, [z]*e) if e else 1+0j
+                if str(r_pow) == str(r_pro):
+                    continue
+
+                self.assertNotIn(z.real, {0.0, -0.0, INF, -INF, NAN})
+                self.assertNotIn(z.imag, {0.0, -0.0, INF, -INF, NAN})
+
+                # We might fail here, because associativity of multiplication
+                # is broken already for floats.
+                # Consider z = 1-1j.  Then z*z*z*z = ((z*z)*z)*z = -4+0j,
+                # while in the algorithm for pow() a diffenent grouping
+                # of operations is used: z**4 = (z*z)*(z*z) = -4-0j.
+                #
+                # Fallback to the generic complex power algorithm.
+                r_pro, r_pro_errno = _testcapi._py_c_pow(z, e)
+                self.assertEqual(r_pro_errno, 0)
+                self.assertClose(r_pow, r_pro)
+                if not isnan(r_pow.real):
+                    self.assertEqual(copysign(1, r_pow.real),
+                                     copysign(1, r_pro.real))
+                else:
+                    self.assertTrue(isnan(r_pro.real))
+                if not isnan(r_pow.imag):
+                    self.assertEqual(copysign(1, r_pow.imag),
+                                     copysign(1, r_pro.imag))
+                else:
+                    self.assertTrue(isnan(r_pro.imag))
+
     def test_boolcontext(self):
         for i in range(100):
             self.assertTrue(complex(random() + 1e-6, random() + 1e-6))

Misc/NEWS.d/next/Core and Builtins/2024-04-17-19-39-30.gh-issue-117999.oCJkud.rst

@@ -0,0 +1,2 @@
+Fixed small nonnegative integer powers for complex numbers with special values
+(``±0.0``, infinities or nan's) in components.  Patch by Sergey B Kirpichev.

Objects/complexobject.c

@@ -156,14 +156,16 @@ _Py_c_pow(Py_complex a, Py_complex b)
     return r;
 }
 
+#define INT_EXP_CUTOFF 100
+
 static Py_complex
 c_powu(Py_complex x, long n)
 {
-    Py_complex r, p;
+    Py_complex p = x, r = n-- ? p : c_1;
     long mask = 1;
-    r = c_1;
-    p = x;
-    while (mask > 0 && n >= mask) {
+    assert(-1 <= n && n <= INT_EXP_CUTOFF);
+    while (n >= mask) {
+        assert(mask>0);
         if (n & mask)
             r = _Py_c_prod(r,p);
         mask <<= 1;
@@ -175,11 +177,10 @@ c_powu(Py_complex x, long n)
 static Py_complex
 c_powi(Py_complex x, long n)
 {
-    if (n > 0)
-        return c_powu(x,n);
-    else
+    if (n < 0)
         return _Py_c_quot(c_1, c_powu(x,-n));
-
+    else
+        return c_powu(x,n);
 }
 
 double
@@ -544,7 +545,7 @@ complex_pow(PyObject *v, PyObject *w, PyObject *z)
     errno = 0;
     // Check whether the exponent has a small integer value, and if so use
     // a faster and more accurate algorithm.
-    if (b.imag == 0.0 && b.real == floor(b.real) && fabs(b.real) <= 100.0) {
+    if (b.imag == 0.0 && b.real == floor(b.real) && fabs(b.real) <= INT_EXP_CUTOFF) {
         p = c_powi(a, (long)b.real);
     }
     else {