python · skirpichev · Sep 20, 2023 · Sep 20, 2023 · Sep 20, 2023 · Sep 20, 2023
diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py
@@ -109,6 +109,8 @@ def test_truediv(self):
                            complex(random(), random()))
 
         self.assertAlmostEqual(complex.__truediv__(2+0j, 1+1j), 1-1j)
+        self.assertRaises(TypeError, operator.truediv, 1j, None)
+        self.assertRaises(TypeError, operator.truediv, None, 1j)
 
         for denom_real, denom_imag in [(0, NAN), (NAN, 0), (NAN, NAN)]:
             z = complex(0, 0) / complex(denom_real, denom_imag)
@@ -140,6 +142,7 @@ def test_floordiv_zero_division(self):
     def test_richcompare(self):
         self.assertIs(complex.__eq__(1+1j, 1<<10000), False)
         self.assertIs(complex.__lt__(1+1j, None), NotImplemented)
+        self.assertIs(complex.__eq__(1+1j, None), NotImplemented)
         self.assertIs(complex.__eq__(1+1j, 1+1j), True)
         self.assertIs(complex.__eq__(1+1j, 2+2j), False)
         self.assertIs(complex.__ne__(1+1j, 1+1j), False)
@@ -162,6 +165,7 @@ def test_richcompare(self):
         self.assertIs(operator.eq(1+1j, 2+2j), False)
         self.assertIs(operator.ne(1+1j, 1+1j), False)
         self.assertIs(operator.ne(1+1j, 2+2j), True)
+        self.assertIs(operator.eq(1+1j, 2.0), False)
 
     def test_richcompare_boundaries(self):
         def check(n, deltas, is_equal, imag = 0.0):
@@ -180,6 +184,27 @@ def check(n, deltas, is_equal, imag = 0.0):
             check(2 ** pow, range(1, 101), lambda delta: False, float(i))
         check(2 ** 53, range(-100, 0), lambda delta: True)
 
+    def test_add(self):
+        self.assertAlmostEqual(1j + 1, complex(+1, 1))
+        self.assertAlmostEqual(1j + (-1), complex(-1, 1))
+        self.assertRaises(OverflowError, operator.add, 1j, 10**1000)
+        self.assertRaises(TypeError, operator.add, 1j, None)
+        self.assertRaises(TypeError, operator.add, None, 1j)
+
+    def test_sub(self):
+        self.assertAlmostEqual(1j - 1, complex(-1, 1))
+        self.assertAlmostEqual(1j - (-1), complex(1, 1))
+        self.assertRaises(OverflowError, operator.sub, 1j, 10**1000)
+        self.assertRaises(TypeError, operator.sub, 1j, None)
+        self.assertRaises(TypeError, operator.sub, None, 1j)
+
+    def test_mul(self):
+        self.assertAlmostEqual(1j * 2, complex(0, 2))
+        self.assertAlmostEqual(1j * (-1), complex(0, -1))
+        self.assertRaises(OverflowError, operator.mul, 1j, 10**1000)
+        self.assertRaises(TypeError, operator.mul, 1j, None)
+        self.assertRaises(TypeError, operator.mul, None, 1j)
+
     def test_mod(self):
         # % is no longer supported on complex numbers
         with self.assertRaises(TypeError):
@@ -212,11 +237,18 @@ def test_divmod_zero_division(self):
     def test_pow(self):
         self.assertAlmostEqual(pow(1+1j, 0+0j), 1.0)
         self.assertAlmostEqual(pow(0+0j, 2+0j), 0.0)
+        self.assertAlmostEqual(pow(0+0j, 2000+0j), 0.0)
+        self.assertAlmostEqual(pow(0, 0+0j), 1.0)
+        self.assertAlmostEqual(pow(-1, 0+0j), 1.0)
         self.assertRaises(ZeroDivisionError, pow, 0+0j, 1j)
+        self.assertRaises(ZeroDivisionError, pow, 0+0j, -1000)
         self.assertAlmostEqual(pow(1j, -1), 1/1j)
         self.assertAlmostEqual(pow(1j, 200), 1)
         self.assertRaises(ValueError, pow, 1+1j, 1+1j, 1+1j)
         self.assertRaises(OverflowError, pow, 1e200+1j, 1e200+1j)
+        self.assertRaises(TypeError, pow, 1j, None)
+        self.assertRaises(TypeError, pow, None, 1j)
+        self.assertAlmostEqual(pow(1j, 0.5), 0.7071067811865476+0.7071067811865475j)
 
         a = 3.33+4.43j
         self.assertEqual(a ** 0j, 1)
@@ -302,10 +334,15 @@ def test_boolcontext(self):
             self.assertTrue(complex(random() + 1e-6, random() + 1e-6))
         self.assertTrue(not complex(0.0, 0.0))
 
+    def test_bool(self):
+        self.assertTrue(1j)
+
     def test_conjugate(self):
         self.assertClose(complex(5.3, 9.8).conjugate(), 5.3-9.8j)
 
     def test_constructor(self):
+        from test.test_capi.test_getargs import Complex
+
         class NS:
             def __init__(self, value): self.value = value
             def __complex__(self): return self.value
@@ -314,6 +351,8 @@ def __complex__(self): return self.value
         self.assertRaises(TypeError, complex, {})
         self.assertRaises(TypeError, complex, NS(1.5))
         self.assertRaises(TypeError, complex, NS(1))
+        self.assertRaises(TypeError, complex, object())
+        self.assertRaises(TypeError, complex, Complex(), object())
 
         self.assertAlmostEqual(complex("1+10j"), 1+10j)
         self.assertAlmostEqual(complex(10), 10+0j)
@@ -359,6 +398,8 @@ def __complex__(self): return self.value
         self.assertAlmostEqual(complex('1e-500'), 0.0 + 0.0j)
         self.assertAlmostEqual(complex('-1e-500j'), 0.0 - 0.0j)
         self.assertAlmostEqual(complex('-1e-500+1e-500j'), -0.0 + 0.0j)
+        self.assertAlmostEqual(complex('1-1j'), 1.0 - 1j)
+        self.assertAlmostEqual(complex('1J'), 1j)
 
         class complex2(complex): pass
         self.assertAlmostEqual(complex(complex2(1+1j)), 1+1j)
@@ -553,6 +594,8 @@ def test_hash(self):
             x /= 3.0    # now check against floating point
             self.assertEqual(hash(x), hash(complex(x, 0.)))
 
+        self.assertEqual(hash(2000005 - 1j), -2)
+
     def test_abs(self):
         nums = [complex(x/3., y/7.) for x in range(-9,9) for y in range(-9,9)]
         for num in nums:
@@ -602,6 +645,12 @@ def test(v, expected, test_fn=self.assertEqual):
         test(complex(-0., 0.),  "(-0+0j)")
         test(complex(-0., -0.), "(-0-0j)")
 
+    def test_pos(self):
+        from test.test_capi.test_getargs import ComplexSubclass
+
+        self.assertEqual(+(1+6j), 1+6j)
+        self.assertEqual(+ComplexSubclass(1, 6), 1+6j)
+
     def test_neg(self):
         self.assertEqual(-(1+6j), -1-6j)
 

diff --git a/Objects/complexobject.c b/Objects/complexobject.c
@@ -131,11 +131,8 @@ _Py_c_pow(Py_complex a, Py_complex b)
 {
     Py_complex r;
     double vabs,len,at,phase;
-    if (b.real == 0. && b.imag == 0.) {
-        r.real = 1.;
-        r.imag = 0.;
-    }
-    else if (a.real == 0. && a.imag == 0.) {
+    assert(!(b.real == 0. && b.imag == 0.));
+    if (a.real == 0. && a.imag == 0.) {
         if (b.imag != 0. || b.real < 0.)
             errno = EDOM;
         r.real = 0.;
@@ -156,14 +153,18 @@ _Py_c_pow(Py_complex a, Py_complex b)
     return r;
 }
 
+#define C_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 && n <= C_EXP_CUTOFF);
+    while (n >= mask) {
+        assert(mask>0);
         if (n & mask)
             r = _Py_c_prod(r,p);
         mask <<= 1;
@@ -408,11 +409,9 @@ complex_hash(PyComplexObject *v)
 {
     Py_uhash_t hashreal, hashimag, combined;
     hashreal = (Py_uhash_t)_Py_HashDouble((PyObject *) v, v->cval.real);
-    if (hashreal == (Py_uhash_t)-1)
-        return -1;
+    assert(hashreal != (Py_uhash_t)-1);
     hashimag = (Py_uhash_t)_Py_HashDouble((PyObject *)v, v->cval.imag);
-    if (hashimag == (Py_uhash_t)-1)
-        return -1;
+    assert(hashimag != (Py_uhash_t)-1);
     /* Note:  if the imaginary part is 0, hashimag is 0 now,
      * so the following returns hashreal unchanged.  This is
      * important because numbers of different types that
@@ -519,7 +518,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) <= C_EXP_CUTOFF) {
         p = c_powi(a, (long)b.real);
     }
     else {
@@ -592,7 +591,7 @@ complex_richcompare(PyObject *v, PyObject *w, int op)
     }
 
     assert(PyComplex_Check(v));
-    TO_COMPLEX(v, i);
+    i = ((PyComplexObject*)v)->cval;
 
     if (PyLong_Check(w)) {
         /* Check for 0.0 imaginary part first to avoid the rich
@@ -618,7 +617,7 @@ complex_richcompare(PyObject *v, PyObject *w, int op)
     else if (PyComplex_Check(w)) {
         Py_complex j;
 
-        TO_COMPLEX(w, j);
+        j = ((PyComplexObject*)w)->cval;
         equal = (i.real == j.real && i.imag == j.imag);
     }
     else {
@@ -857,22 +856,15 @@ complex_subtype_from_string(PyTypeObject *type, PyObject *v)
     PyObject *s_buffer = NULL, *result = NULL;
     Py_ssize_t len;
 
-    if (PyUnicode_Check(v)) {
-        s_buffer = _PyUnicode_TransformDecimalAndSpaceToASCII(v);
-        if (s_buffer == NULL) {
-            return NULL;
-        }
-        assert(PyUnicode_IS_ASCII(s_buffer));
-        /* Simply get a pointer to existing ASCII characters. */
-        s = PyUnicode_AsUTF8AndSize(s_buffer, &len);
-        assert(s != NULL);
-    }
-    else {
-        PyErr_Format(PyExc_TypeError,
-            "complex() argument must be a string or a number, not '%.200s'",
-            Py_TYPE(v)->tp_name);
+    assert(PyUnicode_Check(v));
+    s_buffer = _PyUnicode_TransformDecimalAndSpaceToASCII(v);
+    if (s_buffer == NULL) {
         return NULL;
     }
+    assert(PyUnicode_IS_ASCII(s_buffer));
+    /* Simply get a pointer to existing ASCII characters. */
+    s = PyUnicode_AsUTF8AndSize(s_buffer, &len);
+    assert(s != NULL);
 
     result = _Py_string_to_number_with_underscores(s, len, "complex", v, type,
                                                    complex_from_string_inner);
@@ -948,9 +940,7 @@ complex_new_impl(PyTypeObject *type, PyObject *r, PyObject *i)
                      "complex() first argument must be a string or a number, "
                      "not '%.200s'",
                      Py_TYPE(r)->tp_name);
-        if (own_r) {
-            Py_DECREF(r);
-        }
+        assert(!own_r);
         return NULL;
     }
     if (i != NULL) {
@@ -982,20 +972,17 @@ complex_new_impl(PyTypeObject *type, PyObject *r, PyObject *i)
            value is (properly) of the builtin complex type. */
         cr = ((PyComplexObject*)r)->cval;
         cr_is_complex = 1;
-        if (own_r) {
-            Py_DECREF(r);
-        }
+        assert(own_r);
+        Py_DECREF(r);
     }
     else {
         /* The "real" part really is entirely real, and contributes
            nothing in the imaginary direction.
            Just treat it as a double. */
         tmp = PyNumber_Float(r);
-        if (own_r) {
-            /* r was a newly created complex number, rather
-               than the original "real" argument. */
-            Py_DECREF(r);
-        }
+        /* r was a newly created complex number, rather
+           than the original "real" argument. */
+        assert(!own_r);
         if (tmp == NULL)
             return NULL;
         assert(PyFloat_Check(tmp));