From 4c72a3c85326b96a0bcb7ce8d6ce63aedcfc2cb1 Mon Sep 17 00:00:00 2001
From: Abe <2319792+doyshinda@users.noreply.github.com>
Date: Fri, 6 Mar 2020 18:02:45 -0700
Subject: [PATCH] Implement format code `e` & `E` for floats and ints

---
 tests/snippets/strings.py | 21 +++++++++++++++++++++
 vm/src/format.rs          | 34 ++++++++++++++++++++++------------
 2 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/tests/snippets/strings.py b/tests/snippets/strings.py
index e676d067fc..7471b70050 100644
--- a/tests/snippets/strings.py
+++ b/tests/snippets/strings.py
@@ -450,3 +450,24 @@ assert '{:%}'.format(float('inf')) == 'inf%'
 assert '{:.2%}'.format(float('inf')) == 'inf%'
 with AssertRaises(ValueError, msg='Invalid format specifier'):
     f'{10.0:%3}'
+
+# Test e & E formatting
+assert '{:e}'.format(10) == '1.000000e+01'
+assert '{:.2e}'.format(11) == '1.10e+01'
+assert '{:e}'.format(10.0) == '1.000000e+01'
+assert '{:e}'.format(-10.0) == '-1.000000e+01'
+assert '{:.2e}'.format(10.0) == '1.00e+01'
+assert '{:.2e}'.format(-10.0) == '-1.00e+01'
+assert '{:.2e}'.format(10.1) == '1.01e+01'
+assert '{:.2e}'.format(-10.1) == '-1.01e+01'
+assert '{:.2e}'.format(10.001) == '1.00e+01'
+assert '{:.4e}'.format(100.234) == '1.0023e+02'
+assert '{:.5e}'.format(100.234) == '1.00234e+02'
+assert '{:.2E}'.format(10.0) == '1.00E+01'
+assert '{:.2E}'.format(-10.0) == '-1.00E+01'
+assert '{:e}'.format(float('nan')) == 'nan'
+assert '{:e}'.format(float('-nan')) == 'nan'
+assert '{:E}'.format(float('nan')) == 'NAN'
+assert '{:e}'.format(float('inf')) == 'inf'
+assert '{:e}'.format(float('-inf')) == '-inf'
+assert '{:E}'.format(float('inf')) == 'INF'
diff --git a/vm/src/format.rs b/vm/src/format.rs
index 73d444a8e9..a2bb1ce48f 100644
--- a/vm/src/format.rs
+++ b/vm/src/format.rs
@@ -272,6 +272,16 @@ fn parse_format_spec(text: &str) -> Result<FormatSpec, &'static str> {
     })
 }
 
+// Formats floats into Python style exponent notation, by first formatting in Rust style
+// exponent notation (`1.0000e0`), then convert to Python style (`1.0000e+00`).
+fn format_float_as_exponent(precision: usize, magnitude: f64, separator: &str) -> String {
+    let r_exp = format!("{:.*e}", precision, magnitude);
+    let mut parts = r_exp.splitn(2, 'e');
+    let base = parts.next().unwrap();
+    let exponent = parts.next().unwrap().parse::<i64>().unwrap();
+    format!("{}{}+{:02}", base, separator, exponent)
+}
+
 impl FormatSpec {
     pub fn parse(text: &str) -> Result<FormatSpec, &'static str> {
         parse_format_spec(text)
@@ -370,12 +380,16 @@ impl FormatSpec {
             Some(FormatType::GeneralFormatLower) => {
                 Err("Format code 'g' for object of type 'float' not implemented yet")
             }
-            Some(FormatType::ExponentUpper) => {
-                Err("Format code 'E' for object of type 'float' not implemented yet")
-            }
-            Some(FormatType::ExponentLower) => {
-                Err("Format code 'e' for object of type 'float' not implemented yet")
-            }
+            Some(FormatType::ExponentUpper) => match magnitude {
+                magnitude if magnitude.is_nan() => Ok("NAN".to_owned()),
+                magnitude if magnitude.is_infinite() => Ok("INF".to_owned()),
+                _ => Ok(format_float_as_exponent(precision, magnitude, "E")),
+            },
+            Some(FormatType::ExponentLower) => match magnitude {
+                magnitude if magnitude.is_nan() => Ok("nan".to_owned()),
+                magnitude if magnitude.is_infinite() => Ok("inf".to_owned()),
+                _ => Ok(format_float_as_exponent(precision, magnitude, "e")),
+            },
             Some(FormatType::Percentage) => match magnitude {
                 magnitude if magnitude.is_nan() => Ok("nan%".to_owned()),
                 magnitude if magnitude.is_infinite() => Ok("inf%".to_owned()),
@@ -443,14 +457,10 @@ impl FormatSpec {
             Some(FormatType::GeneralFormatLower) => {
                 Err("Unknown format code 'g' for object of type 'int'")
             }
-            Some(FormatType::ExponentUpper) => {
-                Err("Unknown format code 'E' for object of type 'int'")
-            }
-            Some(FormatType::ExponentLower) => {
-                Err("Unknown format code 'e' for object of type 'int'")
-            }
             Some(FormatType::FixedPointUpper)
             | Some(FormatType::FixedPointLower)
+            | Some(FormatType::ExponentUpper)
+            | Some(FormatType::ExponentLower)
             | Some(FormatType::Percentage) => match num.to_f64() {
                 Some(float) => return self.format_float(float),
                 _ => Err("Unable to convert int to float"),