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

Before, handling of numbers with special values in components (infinities, nans, signed zero) was invalid. Simple example: >>> z = complex(1, -0.0) >>> z*z (1-0j) >>> z**2 (1+0j) Now: >>> z**2 (1-0j)
python · skirpichev · Apr 17, 2024 · May 29, 2024 · May 30, 2024 · May 31, 2024
commit 70d21c765f00adec71ea12913a38d6bd833339c6
diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py
@@ -6,6 +6,9 @@
 
 from random import random
 from math import atan2, isnan, copysign
+from cmath import log, exp, isclose, isnan as cisnan
+from functools import reduce
+from itertools import combinations
 import operator
 
 INF = float("inf")
@@ -330,6 +333,31 @@ 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
+                test = str(r_pow) == str(r_pro)
+                if not test:
+                    # 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.
+                    r_pro = exp(e*log(z))
+                self.assertTrue(test or isclose(r_pow, r_pro))
+                if not cisnan(r_pow):
+                    self.assertEqual(copysign(1, r_pow.real), copysign(1, r_pro.real))
+                    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
@@ -159,10 +159,8 @@ _Py_c_pow(Py_complex a, Py_complex b)
 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) {
         if (n & mask)
             r = _Py_c_prod(r,p);
@@ -175,11 +173,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
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.