Add support for string formatting of floats with ":f" format code

doyshinda · doyshinda · commit f4334db22774 · 2019-12-29T09:33:48.000-07:00
Also adds automatic conversion of int -> float when the ":f" format code is specified. Fixes a f-string formatting style regression introduced in #1625.
diff --git a/parser/src/fstring.rs b/parser/src/fstring.rs
@@ -83,7 +83,7 @@ impl<'a> FStringParser<'a> {
                         }))
                     } else {
                         spec = Some(Box::new(Constant {
-                            value: spec_expression.trim().to_string(),
+                            value: spec_expression.to_string(),
                         }))
                     }
                 }
diff --git a/tests/snippets/strings.py b/tests/snippets/strings.py
@@ -398,3 +398,36 @@ def try_mutate_str():
 # >>> '{{literal}}'.format_map('foo')
 # '{literal}'
 assert '{{literal}}'.format_map('foo') == '{literal}'
+
+# test formatting float values
+assert f'{5:f}' == '5.000000'
+assert f'{-5:f}' == '-5.000000'
+assert f'{5.0:f}' == '5.000000'
+assert f'{-5.0:f}' == '-5.000000'
+assert f'{5:.2f}' == '5.00'
+assert f'{5.0:.2f}' == '5.00'
+assert f'{-5:.2f}' == '-5.00'
+assert f'{-5.0:.2f}' == '-5.00'
+assert f'{5.0:04f}' == '5.000000'
+assert f'{5.1234:+f}' == '+5.123400'
+assert f'{5.1234: f}' == ' 5.123400'
+assert f'{5.1234:-f}' == '5.123400'
+assert f'{-5.1234:-f}' == '-5.123400'
+assert f'{1.0:+}' == '+1.0'
+assert f'--{1.0:f>4}--' == '--f1.0--'
+assert f'--{1.0:f<4}--' == '--1.0f--'
+assert f'--{1.0:d^4}--' == '--1.0d--'
+assert f'--{1.0:d^5}--' == '--d1.0d--'
+assert f'--{1.1:f>6}--' == '--fff1.1--'
+assert '{}'.format(float('nan')) == 'nan'
+assert '{:f}'.format(float('nan')) == 'nan'
+assert '{:f}'.format(float('-nan')) == 'nan'
+assert '{:F}'.format(float('nan')) == 'NAN'
+assert '{}'.format(float('inf')) == 'inf'
+assert '{:f}'.format(float('inf')) == 'inf'
+assert '{:f}'.format(float('-inf')) == '-inf'
+assert '{:F}'.format(float('inf')) == 'INF'
+assert f'{1234567890.1234:,.2f}' == '1,234,567,890.12'
+assert f'{1234567890.1234:_.2f}' == '1_234_567_890.12'
+with AssertRaises(ValueError, msg="Unknown format code 'd' for object of type 'float'"):
+    f'{5.0:04d}'
diff --git a/vm/src/format.rs b/vm/src/format.rs
@@ -1,4 +1,5 @@
 use num_bigint::{BigInt, Sign};
+use num_traits::cast::ToPrimitive;
 use num_traits::Signed;
 use std::cmp;
 use std::str::FromStr;
@@ -281,14 +282,22 @@ impl FormatSpec {
         separator: char,
     ) -> String {
         let mut result = String::new();
-        let mut remaining: usize = magnitude_string.len();
-        for c in magnitude_string.chars() {
+
+        // Don't add separators to the floating decimal point of numbers
+        let parts: Vec<&str> = magnitude_string.split('.').collect();
+        let magnitude_integer_string = parts[0];
+        let mut remaining: usize = magnitude_integer_string.len();
+        for c in magnitude_integer_string.chars() {
             result.push(c);
             remaining -= 1;
             if remaining % interval == 0 && remaining > 0 {
                 result.push(separator);
             }
         }
+        if parts.len() == 2 {
+            result.push('.');
+            result.push_str(parts[1]);
+        }
         result
     }
 
@@ -300,6 +309,7 @@ impl FormatSpec {
             Some(FormatType::HexLower) => 4,
             Some(FormatType::HexUpper) => 4,
             Some(FormatType::Number) => 3,
+            Some(FormatType::FixedPointLower) | Some(FormatType::FixedPointUpper) => 3,
             None => 3,
             _ => panic!("Separators only valid for numbers!"),
         }
@@ -321,8 +331,75 @@ impl FormatSpec {
         }
     }
 
+    pub fn format_float(&self, num: f64) -> Result<String, &'static str> {
+        let precision = self.precision.unwrap_or(6);
+        let magnitude = num.abs();
+        let raw_magnitude_string_result: Result<String, &'static str> = match self.format_type {
+            Some(FormatType::FixedPointUpper) => match magnitude {
+                magnitude if magnitude.is_nan() => Ok("NAN".to_string()),
+                magnitude if magnitude.is_infinite() => Ok("INF".to_string()),
+                _ => Ok(format!("{:.*}", precision, magnitude)),
+            },
+            Some(FormatType::FixedPointLower) => match magnitude {
+                magnitude if magnitude.is_nan() => Ok("nan".to_string()),
+                magnitude if magnitude.is_infinite() => Ok("inf".to_string()),
+                _ => Ok(format!("{:.*}", precision, magnitude)),
+            },
+            Some(FormatType::Decimal) => Err("Unknown format code 'd' for object of type 'float'"),
+            Some(FormatType::Binary) => Err("Unknown format code 'b' for object of type 'float'"),
+            Some(FormatType::Octal) => Err("Unknown format code 'o' for object of type 'float'"),
+            Some(FormatType::HexLower) => Err("Unknown format code 'x' for object of type 'float'"),
+            Some(FormatType::HexUpper) => Err("Unknown format code 'X' for object of type 'float'"),
+            Some(FormatType::String) => Err("Unknown format code 's' for object of type 'float'"),
+            Some(FormatType::Character) => {
+                Err("Unknown format code 'c' for object of type 'float'")
+            }
+            Some(FormatType::Number) => {
+                Err("Format code 'n' for object of type 'float' not implemented yet")
+            }
+            Some(FormatType::GeneralFormatUpper) => {
+                Err("Format code 'G' for object of type 'float' not implemented yet")
+            }
+            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")
+            }
+            None => {
+                match magnitude {
+                    magnitude if magnitude.is_nan() => Ok("nan".to_string()),
+                    magnitude if magnitude.is_infinite() => Ok("inf".to_string()),
+                    // Using the Debug format here to prevent the automatic conversion of floats
+                    // ending in .0 to their integer representation (e.g., 1.0 -> 1)
+                    _ => Ok(format!("{:?}", magnitude)),
+                }
+            }
+        };
+
+        if raw_magnitude_string_result.is_err() {
+            return raw_magnitude_string_result;
+        }
+
+        let magnitude_string = self.add_magnitude_separators(raw_magnitude_string_result.unwrap());
+        let format_sign = self.sign.unwrap_or(FormatSign::Minus);
+        let sign_str = if num.is_sign_negative() && !num.is_nan() {
+            "-"
+        } else {
+            match format_sign {
+                FormatSign::Plus => "+",
+                FormatSign::Minus => "",
+                FormatSign::MinusOrSpace => " ",
+            }
+        };
+
+        self.format_sign_and_align(magnitude_string, sign_str)
+    }
+
     pub fn format_int(&self, num: &BigInt) -> Result<String, &'static str> {
-        let fill_char = self.fill.unwrap_or(' ');
         let magnitude = num.abs();
         let prefix = if self.alternate_form {
             match self.format_type {
@@ -360,11 +437,11 @@ impl FormatSpec {
             Some(FormatType::ExponentLower) => {
                 Err("Unknown format code 'e' for object of type 'int'")
             }
-            Some(FormatType::FixedPointUpper) => {
-                Err("Unknown format code 'F' for object of type 'int'")
-            }
-            Some(FormatType::FixedPointLower) => {
-                Err("Unknown format code 'f' for object of type 'int'")
+            Some(FormatType::FixedPointUpper) | Some(FormatType::FixedPointLower) => {
+                match num.to_f64() {
+                    Some(float) => return self.format_float(float),
+                    _ => Err("Unable to convert int to float"),
+                }
             }
             None => Ok(magnitude.to_str_radix(10)),
         };
@@ -376,10 +453,6 @@ impl FormatSpec {
             prefix,
             self.add_magnitude_separators(raw_magnitude_string_result.unwrap())
         );
-        let align = self.align.unwrap_or(FormatAlign::Right);
-
-        // Use the byte length as the string length since we're in ascii
-        let num_chars = magnitude_string.len();
 
         let format_sign = self.sign.unwrap_or(FormatSign::Minus);
         let sign_str = match num.sign() {
@@ -391,6 +464,19 @@ impl FormatSpec {
             },
         };
 
+        self.format_sign_and_align(magnitude_string, sign_str)
+    }
+
+    fn format_sign_and_align(
+        &self,
+        magnitude_string: String,
+        sign_str: &str,
+    ) -> Result<String, &'static str> {
+        let align = self.align.unwrap_or(FormatAlign::Right);
+
+        // Use the byte length as the string length since we're in ascii
+        let num_chars = magnitude_string.len();
+        let fill_char = self.fill.unwrap_or(' ');
         let fill_chars_needed: i32 = self.width.map_or(0, |w| {
             cmp::max(0, (w as i32) - (num_chars as i32) - (sign_str.len() as i32))
         });
diff --git a/vm/src/obj/objfloat.rs b/vm/src/obj/objfloat.rs
@@ -7,6 +7,7 @@ use super::objbytes;
 use super::objint::{self, PyInt, PyIntRef};
 use super::objstr::{PyString, PyStringRef};
 use super::objtype::{self, PyClassRef};
+use crate::format::FormatSpec;
 use crate::function::{OptionalArg, OptionalOption};
 use crate::pyhash;
 use crate::pyobject::{
@@ -209,6 +210,15 @@ impl PyFloat {
         vm.ctx.new_bool(result)
     }
 
+    #[pymethod(name = "__format__")]
+    fn format(&self, spec: PyStringRef, vm: &VirtualMachine) -> PyResult<String> {
+        let format_spec = FormatSpec::parse(spec.as_str());
+        match format_spec.format_float(self.value) {
+            Ok(string) => Ok(string),
+            Err(err) => Err(vm.new_value_error(err.to_string())),
+        }
+    }
+
     #[pymethod(name = "__eq__")]
     fn eq(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
         self.cmp(other, |a, b| a == b, |a, b| int_eq(a, b), vm)

Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,7 @@ impl<'a> FStringParser<'a> {`
`83`	`83`	`}))`
`84`	`84`	`} else {`
`85`	`85`	`spec = Some(Box::new(Constant {`
`86`		`- value: spec_expression.trim().to_string(),`
	`86`	`+ value: spec_expression.to_string(),`
`87`	`87`	`}))`
`88`	`88`	`}`
`89`	`89`	`}`