repair compile time evaluation of complex numbers

3 files changed, 24 insertions, 10 deletions

diff --git a/Cython/Compiler/ExprNodes.py b/Cython/Compiler/ExprNodes.py
index e364c777f..ee8ad1d19 100644
--- a/Cython/Compiler/ExprNodes.py
+++ b/Cython/Compiler/ExprNodes.py
@@ -1649,15 +1649,15 @@ class IdentifierStringNode(StringNode):
 class ImagNode(AtomicExprNode):
     #  Imaginary number literal
     #
-    #  value   float    imaginary part
+    #  value   string    imaginary part (float value)
 
     type = PyrexTypes.c_double_complex_type
 
     def calculate_constant_result(self):
-        self.constant_result = complex(0.0, self.value)
+        self.constant_result = complex(0.0, float(self.value))
 
     def compile_time_value(self, denv):
-        return complex(0.0, self.value)
+        return complex(0.0, float(self.value))
 
     def analyse_types(self, env):
         self.type.create_declaration_utility_code(env)
diff --git a/Cython/Compiler/Parsing.py b/Cython/Compiler/Parsing.py
index a440ee4c8..17b69d39b 100644
--- a/Cython/Compiler/Parsing.py
+++ b/Cython/Compiler/Parsing.py
@@ -771,6 +771,15 @@ def wrap_compile_time_constant(pos, value):
         return ExprNodes.IntNode(pos, value=rep, constant_result=value)
     elif isinstance(value, float):
         return ExprNodes.FloatNode(pos, value=rep, constant_result=value)
+    elif isinstance(value, complex):
+        node = ExprNodes.ImagNode(pos, value=repr(value.imag), constant_result=complex(0.0, value.imag))
+        if value.real:
+            # FIXME: should we care about -0.0 ?
+            # probably not worth using the '-' operator for negative imag values
+            node = ExprNodes.binop_node(
+                pos, '+', ExprNodes.FloatNode(pos, value=repr(value.real), constant_result=value.real), node,
+                constant_result=value)
+        return node
     elif isinstance(value, _unicode):
         return ExprNodes.UnicodeNode(pos, value=EncodedString(value))
     elif isinstance(value, _bytes):
diff --git a/tests/run/complex_numbers_T305.pyx b/tests/run/complex_numbers_T305.pyx
index 9d719f609..310b7233e 100644
--- a/tests/run/complex_numbers_T305.pyx
+++ b/tests/run/complex_numbers_T305.pyx
@@ -4,6 +4,8 @@ from cpython.object cimport Py_EQ, Py_NE
 
 cimport cython
 
+DEF C21 = 2-1j
+
 
 cdef class Complex3j:
     """
@@ -137,15 +139,17 @@ def test_coercion(int a, float b, double c, float complex d, double complex e):
 def test_compare(double complex a, double complex b):
     """
     >>> test_compare(3, 3)
-    (True, False, False, False, False, True)
+    (True, False, False, False, False, True, False)
     >>> test_compare(3j, 3j)
-    (True, False, True, True, True, False)
+    (True, False, True, True, True, False, False)
     >>> test_compare(3j, 4j)
-    (False, True, True, False, True, True)
+    (False, True, True, False, True, True, False)
     >>> test_compare(3, 4)
-    (False, True, False, False, False, True)
+    (False, True, False, False, False, True, False)
+    >>> test_compare(2-1j, 4)
+    (False, True, False, False, False, True, True)
     """
-    return a == b, a != b, a == 3j, 3j == b, a == Complex3j(), Complex3j() != b
+    return a == b, a != b, a == 3j, 3j == b, a == Complex3j(), Complex3j() != b, a == C21
 
 
 def test_compare_coerce(double complex a, int b):
@@ -165,9 +169,10 @@ def test_compare_coerce(double complex a, int b):
 def test_literal():
     """
     >>> test_literal()
-    (5j, (1-2.5j))
+    (5j, (1-2.5j), (2-1j))
     """
-    return 5j, 1-2.5j
+    return 5j, 1-2.5j, C21
+
 
 def test_real_imag(double complex z):
     """