gh-117999: use generic algorithm in complex_pow() if base has special components

Lib/test/test_complex.py

@@ -1,3 +1,4 @@
+import cmath
 import unittest
 import sys
 from test import support
@@ -7,7 +8,9 @@
 
 from random import random
 from math import isnan, copysign
+from itertools import combinations_with_replacement
 import operator
+import _testcapi
 
 INF = float("inf")
 NAN = float("nan")
@@ -445,6 +448,36 @@ 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_with_replacement([1, -1, 0.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:
+                with self.subTest(value=z, exponent=e):
+                    if cmath.isfinite(z) and z.real and z.imag:
+                        continue
+                    try:
+                        r_pow = z**e
+                    except OverflowError:
+                        continue
+                    # Use the generic complex power algorithm.
+                    r_pro, r_pro_errno = _testcapi._py_c_pow(z, e)
+                    self.assertEqual(r_pro_errno, 0)
+                    if isnan(r_pow.real):
+                        self.assertTrue(isnan(r_pro.real))
+                    else:
+                        self.assertEqual(copysign(1, r_pow.real),
+                                         copysign(1, r_pro.real))
+                    if isnan(r_pow.imag):
+                        self.assertTrue(isnan(r_pro.imag))
+                    else:
+                        self.assertEqual(copysign(1, r_pow.imag),
+                                         copysign(1, 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/Library/2024-08-24-10-46-35.gh-issue-117999.5K_BiA.rst

@@ -0,0 +1,2 @@
+Use a single algorithm for complex exponentiation (the case where the exponent
+is a small integer was previously handled separately).  Patch by Sergey B Kirpichev.

Objects/complexobject.c

@@ -338,14 +338,20 @@ _Py_c_pow(Py_complex a, Py_complex b)
     return r;
 }
 
+/* Switch to exponentiation by squaring if integer exponent less that this. */
+#define INT_EXP_CUTOFF 100
+
 static Py_complex
 c_powu(Py_complex x, long n)
 {
     Py_complex r, p;
     long mask = 1;
     r = c_1;
     p = x;
-    while (mask > 0 && n >= mask) {
+    assert(0 <= n);
+    assert(n <= INT_EXP_CUTOFF);
+    while (n >= mask) {
+        assert(mask > 0);
         if (n & mask)
             r = _Py_c_prod(r,p);
         mask <<= 1;
@@ -735,7 +741,11 @@ 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) {
+    // Fallback on the generic code if the base has special
+    // components (zeros or infinities).
+    if (b.imag == 0.0 && b.real == floor(b.real) && fabs(b.real) <= INT_EXP_CUTOFF
+        && isfinite(a.real) && a.real && isfinite(a.imag) && a.imag)
+    {
         p = c_powi(a, (long)b.real);
         _Py_ADJUST_ERANGE2(p.real, p.imag);
     }