From 0172292762649db91f588107b2163ab6449dc7ca Mon Sep 17 00:00:00 2001
From: Jim Mussared <jim.mussared@gmail.com>
Date: Sat, 7 Aug 2021 02:16:59 +1000
Subject: [PATCH] py/parsenum: Support parsing complex numbers of the form
 "a+bj".

To conform with CPython.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>
---
 py/parsenum.c           | 37 ++++++++++++++++++++++++++++++++-----
 tests/float/complex1.py | 12 ++++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/py/parsenum.c b/py/parsenum.c
index 1cfe842577..848a7732d9 100644
--- a/py/parsenum.c
+++ b/py/parsenum.c
@@ -166,6 +166,12 @@ value_error:
     }
 }
 
+enum {
+    REAL_IMAG_STATE_START = 0,
+    REAL_IMAG_STATE_HAVE_REAL = 1,
+    REAL_IMAG_STATE_HAVE_IMAG = 2,
+};
+
 typedef enum {
     PARSE_DEC_IN_INTG,
     PARSE_DEC_IN_FRAC,
@@ -196,7 +202,12 @@ mp_obj_t mp_parse_num_decimal(const char *str, size_t len, bool allow_imag, bool
     const char *top = str + len;
     mp_float_t dec_val = 0;
     bool dec_neg = false;
-    bool imag = false;
+    unsigned int real_imag_state = REAL_IMAG_STATE_START;
+
+    #if MICROPY_PY_BUILTINS_COMPLEX
+    mp_float_t dec_real = 0;
+parse_start:
+    #endif
 
     // skip leading space
     for (; str < top && unichar_isspace(*str); str++) {
@@ -281,7 +292,7 @@ mp_obj_t mp_parse_num_decimal(const char *str, size_t len, bool allow_imag, bool
                     goto value_error;
                 }
             } else if (allow_imag && (dig | 0x20) == 'j') {
-                imag = true;
+                real_imag_state |= REAL_IMAG_STATE_HAVE_IMAG;
                 break;
             } else if (dig == '_') {
                 continue;
@@ -332,18 +343,34 @@ mp_obj_t mp_parse_num_decimal(const char *str, size_t len, bool allow_imag, bool
 
     // check we reached the end of the string
     if (str != top) {
+        #if MICROPY_PY_BUILTINS_COMPLEX
+        if (force_complex && real_imag_state == REAL_IMAG_STATE_START) {
+            // If we've only seen a real so far, keep parsing for the imaginary part.
+            dec_real = dec_val;
+            dec_val = 0;
+            real_imag_state |= REAL_IMAG_STATE_HAVE_REAL;
+            goto parse_start;
+        }
+        #endif
         goto value_error;
     }
 
+    #if MICROPY_PY_BUILTINS_COMPLEX
+    if (real_imag_state == REAL_IMAG_STATE_HAVE_REAL) {
+        // We're on the second part, but didn't get the expected imaginary number.
+        goto value_error;
+    }
+    #endif
+
     // return the object
     #if MICROPY_PY_BUILTINS_COMPLEX
-    if (imag) {
-        return mp_obj_new_complex(0, dec_val);
+    if (real_imag_state != REAL_IMAG_STATE_START) {
+        return mp_obj_new_complex(dec_real, dec_val);
     } else if (force_complex) {
         return mp_obj_new_complex(dec_val, 0);
     }
     #else
-    if (imag || force_complex) {
+    if (real_imag_state != REAL_IMAG_STATE_START || force_complex) {
         raise_exc(mp_obj_new_exception_msg(&mp_type_ValueError, MP_ERROR_TEXT("complex values not supported")), lex);
     }
     #endif
diff --git a/tests/float/complex1.py b/tests/float/complex1.py
index 139bb0c509..feede0eab3 100644
--- a/tests/float/complex1.py
+++ b/tests/float/complex1.py
@@ -7,6 +7,11 @@ print(complex(1.2j))
 print(complex("1"))
 print(complex("1.2"))
 print(complex("1.2j"))
+print(complex("1+2j"))
+print(complex("-1-2j"))
+print(complex("+1-2j"))
+print(complex(" -1-2j "))
+print(complex(" +1-2j "))
 print(complex(1, 2))
 print(complex(1j, 2j))
 
@@ -72,6 +77,13 @@ print(float("-nan") * 1j)
 print(float("inf") * (1 + 1j))
 print(float("-inf") * (1 + 1j))
 
+# malformed complex strings
+for test in ("1+2", "1j+2", "1+2j+3", "1+2+3j", "1 + 2j"):
+    try:
+        complex(test)
+    except ValueError:
+        print("ValueError", test)
+
 # can't assign to attributes
 try:
     (1j).imag = 0