From ec5d168de20977c91500012b44845da190742307 Mon Sep 17 00:00:00 2001
From: Youngchan Lee <youngchan1115@gmail.com>
Date: Tue, 18 Aug 2020 02:19:58 +0900
Subject: [PATCH] Implement objcomplex::parse_str() to replace
 Complex64::from_str()

---
 tests/snippets/builtin_complex.py | 45 +++++++++++++++++++++++++++++++
 vm/src/obj/objcomplex.rs          | 45 ++++++++++++++++++++++++++++---
 2 files changed, 86 insertions(+), 4 deletions(-)

diff --git a/tests/snippets/builtin_complex.py b/tests/snippets/builtin_complex.py
index 7d65abf16..58c9b5384 100644
--- a/tests/snippets/builtin_complex.py
+++ b/tests/snippets/builtin_complex.py
@@ -164,3 +164,48 @@ assert complex("5-2j") == 5 - 2j
 assert complex("-2j") == -2j
 assert_raises(TypeError, lambda: complex("5+2j", 1))
 assert_raises(ValueError, lambda: complex("abc"))
+
+assert complex("1+10j") == 1+10j
+assert complex(10) == 10+0j
+assert complex(10.0) == 10+0j
+assert complex(10) == 10+0j
+assert complex(10+0j) == 10+0j
+assert complex(1, 10) == 1+10j
+assert complex(1, 10) == 1+10j
+assert complex(1, 10.0) == 1+10j
+assert complex(1, 10) == 1+10j
+assert complex(1, 10) == 1+10j
+assert complex(1, 10.0) == 1+10j
+assert complex(1.0, 10) == 1+10j
+assert complex(1.0, 10) == 1+10j
+assert complex(1.0, 10.0) == 1+10j
+assert complex(3.14+0j) == 3.14+0j
+assert complex(3.14) == 3.14+0j
+assert complex(314) == 314.0+0j
+assert complex(314) == 314.0+0j
+assert complex(3.14+0j, 0j) == 3.14+0j
+assert complex(3.14, 0.0) == 3.14+0j
+assert complex(314, 0) == 314.0+0j
+assert complex(314, 0) == 314.0+0j
+assert complex(0j, 3.14j) == -3.14+0j
+assert complex(0.0, 3.14j) == -3.14+0j
+assert complex(0j, 3.14) == 3.14j
+assert complex(0.0, 3.14) == 3.14j
+assert complex("1") == 1+0j
+assert complex("1j") == 1j
+assert complex() == 0
+assert complex("-1") == -1
+assert complex("+1") == +1
+assert complex("(1+2j)") == 1+2j
+assert complex("(1.3+2.2j)") == 1.3+2.2j
+assert complex("3.14+1J") == 3.14+1j
+assert complex(" ( +3.14-6J )") == 3.14-6j
+assert complex(" ( +3.14-J )") == 3.14-1j
+assert complex(" ( +3.14+j )") == 3.14+1j
+assert complex("J") == 1j
+assert complex("( j )") == 1j
+assert complex("+J") == 1j
+assert complex("( -j)") == -1j
+assert complex('1e-500') == 0.0 + 0.0j
+assert complex('-1e-500j') == 0.0 - 0.0j
+assert complex('-1e-500+1e-500j') == -0.0 + 0.0j
diff --git a/vm/src/obj/objcomplex.rs b/vm/src/obj/objcomplex.rs
index b376e6661..6ce41459e 100644
--- a/vm/src/obj/objcomplex.rs
+++ b/vm/src/obj/objcomplex.rs
@@ -1,6 +1,5 @@
 use num_complex::Complex64;
 use num_traits::Zero;
-use std::str::FromStr;
 
 use super::objfloat;
 use super::objstr::PyString;
@@ -10,7 +9,7 @@ use crate::pyobject::{
     PyObjectRef, PyRef, PyResult, PyValue, TypeProtocol,
 };
 use crate::vm::VirtualMachine;
-use rustpython_common::hash;
+use rustpython_common::{float_ops, hash};
 
 /// Create a complex number from a real part and an optional imaginary part.
 ///
@@ -254,8 +253,9 @@ impl PyComplex {
                             "complex() can't take second arg if first is a string".to_owned(),
                         ));
                     }
-                    let value = Complex64::from_str(s.borrow_value())
-                        .map_err(|err| vm.new_value_error(err.to_string()))?;
+                    let value = parse_str(s.borrow_value().trim()).ok_or_else(|| {
+                        vm.new_value_error("complex() arg is a malformed string".to_owned())
+                    })?;
                     return Self::from(value).into_ref_with_type(vm, cls);
                 } else {
                     return Err(vm.new_type_error(format!(
@@ -308,3 +308,40 @@ struct ComplexArgs {
     #[pyarg(positional_or_keyword, default = "None")]
     imag: Option<PyObjectRef>,
 }
+
+fn parse_str(s: &str) -> Option<Complex64> {
+    // Handle parentheses
+    let s = match s.strip_prefix('(') {
+        None => s,
+        Some(s) => match s.strip_suffix(')') {
+            None => return None,
+            Some(s) => s.trim(),
+        },
+    };
+
+    let value = match s.strip_suffix(|c| c == 'j' || c == 'J') {
+        None => Complex64::new(float_ops::parse_str(s)?, 0.0),
+        Some(mut s) => {
+            let mut real = 0.0;
+            // Find the central +/- operator. If it exists, parse the real part.
+            for (i, w) in s.as_bytes().windows(2).enumerate() {
+                if (w[1] == b'+' || w[1] == b'-') && !(w[0] == b'e' || w[0] == b'E') {
+                    real = float_ops::parse_str(&s[..=i])?;
+                    s = &s[i + 1..];
+                    break;
+                }
+            }
+
+            let imag = match s {
+                // "j", "+j"
+                "" | "+" => 1.0,
+                // "-j"
+                "-" => -1.0,
+                s => float_ops::parse_str(s)?,
+            };
+
+            Complex64::new(real, imag)
+        }
+    };
+    Some(value)
+}