Implement float literal formatting (#3184)

crates/ruff_python_formatter/src/format/expr.rs

@@ -13,7 +13,7 @@ use crate::cst::{
     Arguments, Boolop, Cmpop, Comprehension, Expr, ExprKind, Keyword, Operator, Unaryop,
 };
 use crate::format::helpers::{is_self_closing, is_simple_power, is_simple_slice};
-use crate::format::numbers::int_literal;
+use crate::format::numbers::{float_literal, int_literal};
 use crate::format::strings::string_literal;
 use crate::shared_traits::AsFormat;
 use crate::trivia::{Parenthesize, Relationship, TriviaKind};
@@ -654,6 +654,7 @@ fn format_constant(
             }
         }
         Constant::Int(_) => write!(f, [int_literal(Range::from_located(expr))])?,
+        Constant::Float(_) => write!(f, [float_literal(Range::from_located(expr))])?,
         Constant::Str(_) | Constant::Bytes(_) => write!(f, [string_literal(expr)])?,
         _ => write!(f, [literal(Range::from_located(expr))])?,
     }

crates/ruff_python_formatter/src/format/numbers.rs

@@ -1,11 +1,122 @@
 use ruff_formatter::prelude::*;
 use ruff_formatter::{write, Format};
 use ruff_text_size::TextSize;
+use rustpython_parser::ast::Location;
 
 use crate::builders::literal;
 use crate::context::ASTFormatContext;
 use crate::core::types::Range;
 
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+struct FloatAtom {
+    range: Range,
+}
+
+impl Format<ASTFormatContext<'_>> for FloatAtom {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
+        let (source, start, end) = f.context().locator().slice(self.range);
+
+        if let Some(dot_index) = source[start..end].find('.') {
+            let integer = &source[start..start + dot_index];
+            let fractional = &source[start + dot_index + 1..end];
+
+            if integer.is_empty() {
+                write!(f, [text("0")])?;
+            } else {
+                write!(
+                    f,
+                    [literal(Range::new(
+                        self.range.location,
+                        Location::new(
+                            self.range.location.row(),
+                            self.range.location.column() + dot_index
+                        ),
+                    ))]
+                )?;
+            }
+
+            write!(f, [text(".")])?;
+
+            if fractional.is_empty() {
+                write!(f, [text("0")])?;
+            } else {
+                write!(
+                    f,
+                    [literal(Range::new(
+                        Location::new(
+                            self.range.location.row(),
+                            self.range.location.column() + dot_index + 1
+                        ),
+                        self.range.end_location
+                    ))]
+                )?;
+            }
+        } else {
+            write!(f, [literal(self.range)])?;
+        }
+
+        Ok(())
+    }
+}
+
+#[inline]
+const fn float_atom(range: Range) -> FloatAtom {
+    FloatAtom { range }
+}
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub struct FloatLiteral {
+    range: Range,
+}
+
+impl Format<ASTFormatContext<'_>> for FloatLiteral {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
+        let (source, start, end) = f.context().locator().slice(self.range);
+
+        // Scientific notation
+        if let Some(exponent_index) = source[start..end]
+            .find('e')
+            .or_else(|| source[start..end].find('E'))
+        {
+            // Write the base.
+            write!(
+                f,
+                [float_atom(Range::new(
+                    self.range.location,
+                    Location::new(
+                        self.range.location.row(),
+                        self.range.location.column() + exponent_index
+                    ),
+                ))]
+            )?;
+
+            write!(f, [text("e")])?;
+
+            // Write the exponent, omitting the sign if it's positive.
+            let plus = source[start + exponent_index + 1..end].starts_with('+');
+            write!(
+                f,
+                [literal(Range::new(
+                    Location::new(
+                        self.range.location.row(),
+                        self.range.location.column() + exponent_index + 1 + usize::from(plus)
+                    ),
+                    self.range.end_location
+                ))]
+            )?;
+        } else {
+            write!(f, [float_atom(self.range)])?;
+        }
+
+        Ok(())
+    }
+}
+
+#[inline]
+pub const fn float_literal(range: Range) -> FloatLiteral {
+    FloatLiteral { range }
+}
+
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub struct IntLiteral {
     range: Range,