add proper underscore handling to float parsing

RustPython · youknowone · Jul 28, 2024 · Jul 18, 2024 · Jul 18, 2024 · Jul 18, 2024
commit 9310bda1bad5fca7f8f4b571cc05507853049e11
diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py
@@ -577,7 +577,7 @@ def test_constructor_negative_nans_from_string(self):
         self.assertEqual(copysign(1., complex("-nan-nanj").imag), -1.)
 
     # TODO: RUSTPYTHON
-    @unittest.expectedFailure
+    # @unittest.expectedFailure
     def test_underscores(self):
         # check underscores
         for lit in VALID_UNDERSCORE_LITERALS:

diff --git a/Lib/test/test_float.py b/Lib/test/test_float.py
@@ -73,7 +73,7 @@ def test_noargs(self):
         self.assertEqual(float(), 0.0)
 
     # TODO: RUSTPYTHON
-    @unittest.expectedFailure
+    #@unittest.expectedFailure
     def test_underscores(self):
         for lit in VALID_UNDERSCORE_LITERALS:
             if not any(ch in lit for ch in 'jJxXoObB'):

diff --git a/vm/src/builtins/complex.rs b/vm/src/builtins/complex.rs
@@ -545,13 +545,20 @@ fn parse_str(s: &str) -> Option<Complex64> {
     };
 
     let value = match s.strip_suffix(|c| c == 'j' || c == 'J') {
-        None => Complex64::new(crate::literal::float::parse_str(s)?, 0.0),
+        None => {
+            let stripped = float::float_strip_separators(s.as_bytes())?;
+            Some(Complex64::new(
+                crate::literal::float::parse_bytes(&stripped)?,
+                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 = crate::literal::float::parse_str(&s[..=i])?;
+                    let stripped = float::float_strip_separators(s[..=i].as_bytes())?;
+                    real = crate::literal::float::parse_bytes(&stripped)?;
                     s = &s[i + 1..];
                     break;
                 }
@@ -562,11 +569,13 @@ fn parse_str(s: &str) -> Option<Complex64> {
                 "" | "+" => 1.0,
                 // "-j"
                 "-" => -1.0,
-                s => crate::literal::float::parse_str(s)?,
+                s => {
+                    let stripped = float::float_strip_separators(s.as_bytes())?;
+                    crate::literal::float::parse_bytes(&stripped)?
+                }
             };
-
-            Complex64::new(real, imag)
+            Some(Complex64::new(real, imag))
         }
     };
-    Some(value)
+    value
 }
diff --git a/vm/src/builtins/float.rs b/vm/src/builtins/float.rs
@@ -159,6 +159,30 @@ impl Constructor for PyFloat {
     }
 }
 
+pub fn float_strip_separators(b: &[u8]) -> Option<Vec<u8>> {
+    let mut prev = b'\0';
+    let mut dup = Vec::<u8>::new();
+    for p in b {
+        if *p == b'_' {
+            if !prev.is_ascii_digit() {
+                return None;
+            }
+        } else {
+            dup.push(*p);
+            if prev == b'_' && !p.is_ascii_digit() {
+                return None;
+            }
+        }
+        prev = *p;
+    }
+
+    if prev == b'_' {
+        return None;
+    }
+
+    Some(dup)
+}
+
 fn float_from_string(val: PyObjectRef, vm: &VirtualMachine) -> PyResult<f64> {
     let (bytearray, buffer, buffer_lock);
     let b = if let Some(s) = val.payload_if_subclass::<PyStr>(vm) {
@@ -178,11 +202,19 @@ fn float_from_string(val: PyObjectRef, vm: &VirtualMachine) -> PyResult<f64> {
             val.class().name()
         )));
     };
-    crate::literal::float::parse_bytes(b).ok_or_else(|| {
-        val.repr(vm)
-            .map(|repr| vm.new_value_error(format!("could not convert string to float: {repr}")))
-            .unwrap_or_else(|e| e)
-    })
+
+    let err = val
+        .repr(vm)
+        .map(|repr| vm.new_value_error(format!("could not convert string to float: {repr}")))
+        .unwrap_or_else(|e| e);
+
+    if !b.contains(&b'_') {
+        crate::literal::float::parse_bytes(b).ok_or(err)
+    } else if let Some(dup) = float_strip_separators(b) {
+        crate::literal::float::parse_bytes(&dup).ok_or(err)
+    } else {
+        Err(err)
+    }
 }
 
 #[pyclass(