python · skirpichev · Apr 17, 2024 · May 29, 2024 · May 30, 2024 · May 31, 2024
diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py
@@ -6,7 +6,10 @@
 
 from random import random
 from math import isnan, copysign
+from functools import reduce
+from itertools import combinations_with_replacement
 import operator
+import _testcapi
 
 INF = float("inf")
 NAN = float("nan")
@@ -351,6 +354,47 @@ 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):
+                    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 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))

diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-04-17-19-39-30.gh-issue-117999.oCJkud.rst b/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.
diff --git a/Objects/complexobject.c b/Objects/complexobject.c
@@ -156,14 +156,17 @@ _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);
+    assert(n < INT_EXP_CUTOFF);
+    while (n >= mask) {
+        assert(mask>0);
         if (n & mask)
             r = _Py_c_prod(r,p);
         mask <<= 1;
@@ -175,11 +178,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 +546,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 {
Original file line number	Diff line number	Diff line change
		@@ -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.