py/parsenum: Support parsing complex numbers of the form "a+bj".

To conform with CPython.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>
This commit is contained in:
Jim Mussared 2021-08-07 02:16:59 +10:00 committed by Damien George
parent 7861eddd0f
commit 0172292762
2 changed files with 44 additions and 5 deletions

View File

@ -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

View File

@ -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