Format numeric constants (#5972)

Co-authored-by: Micha Reiser <micha@reiser.io>

crates/ruff_python_formatter/src/expression/expr_constant.rs

@@ -1,14 +1,14 @@
 use ruff_text_size::{TextLen, TextRange};
 use rustpython_parser::ast::{Constant, ExprConstant, Ranged};
 
-use ruff_formatter::write;
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::str::is_implicit_concatenation;
 
+use crate::expression::number::{FormatComplex, FormatFloat, FormatInt};
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::expression::string::{FormatString, StringPrefix, StringQuotes};
 use crate::prelude::*;
-use crate::{not_yet_implemented_custom_text, verbatim_text, FormatNodeRule};
+use crate::{not_yet_implemented_custom_text, FormatNodeRule};
 
 #[derive(Default)]
 pub struct FormatExprConstant;
@@ -28,9 +28,9 @@ impl FormatNodeRule<ExprConstant> for FormatExprConstant {
                 true => text("True").fmt(f),
                 false => text("False").fmt(f),
             },
-            Constant::Int(_) | Constant::Float(_) | Constant::Complex { .. } => {
+            Constant::Int(_) => FormatInt::new(item).fmt(f),
-                write!(f, [verbatim_text(item)])
+            Constant::Float(_) => FormatFloat::new(item).fmt(f),
-            }
+            Constant::Complex { .. } => FormatComplex::new(item).fmt(f),
             Constant::Str(_) => FormatString::new(item).fmt(f),
             Constant::Bytes(_) => {
                 not_yet_implemented_custom_text(r#"b"NOT_YET_IMPLEMENTED_BYTE_STRING""#).fmt(f)

crates/ruff_python_formatter/src/expression/mod.rs

@@ -43,6 +43,7 @@ pub(crate) mod expr_tuple;
 pub(crate) mod expr_unary_op;
 pub(crate) mod expr_yield;
 pub(crate) mod expr_yield_from;
+pub(crate) mod number;
 pub(crate) mod parentheses;
 pub(crate) mod string;
 

crates/ruff_python_formatter/src/expression/number.rs

@@ -0,0 +1,200 @@
+use std::borrow::Cow;
+
+use ruff_text_size::TextSize;
+use rustpython_parser::ast::{ExprConstant, Ranged};
+
+use crate::prelude::*;
+
+pub(super) struct FormatInt<'a> {
+    constant: &'a ExprConstant,
+}
+
+impl<'a> FormatInt<'a> {
+    pub(super) fn new(constant: &'a ExprConstant) -> Self {
+        debug_assert!(constant.value.is_int());
+        Self { constant }
+    }
+}
+
+impl Format<PyFormatContext<'_>> for FormatInt<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let range = self.constant.range();
+        let content = f.context().locator().slice(range);
+
+        let normalized = normalize_integer(content);
+
+        match normalized {
+            Cow::Borrowed(_) => source_text_slice(range, ContainsNewlines::No).fmt(f),
+            Cow::Owned(normalized) => dynamic_text(&normalized, Some(range.start())).fmt(f),
+        }
+    }
+}
+
+pub(super) struct FormatFloat<'a> {
+    constant: &'a ExprConstant,
+}
+
+impl<'a> FormatFloat<'a> {
+    pub(super) fn new(constant: &'a ExprConstant) -> Self {
+        debug_assert!(constant.value.is_float());
+        Self { constant }
+    }
+}
+
+impl Format<PyFormatContext<'_>> for FormatFloat<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let range = self.constant.range();
+        let content = f.context().locator().slice(range);
+
+        let normalized = normalize_floating_number(content);
+
+        match normalized {
+            Cow::Borrowed(_) => source_text_slice(range, ContainsNewlines::No).fmt(f),
+            Cow::Owned(normalized) => dynamic_text(&normalized, Some(range.start())).fmt(f),
+        }
+    }
+}
+
+pub(super) struct FormatComplex<'a> {
+    constant: &'a ExprConstant,
+}
+
+impl<'a> FormatComplex<'a> {
+    pub(super) fn new(constant: &'a ExprConstant) -> Self {
+        debug_assert!(constant.value.is_complex());
+        Self { constant }
+    }
+}
+
+impl Format<PyFormatContext<'_>> for FormatComplex<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let range = self.constant.range();
+        let content = f.context().locator().slice(range);
+
+        let normalized = normalize_floating_number(content.trim_end_matches(['j', 'J']));
+
+        match normalized {
+            Cow::Borrowed(_) => {
+                source_text_slice(range.sub_end(TextSize::from(1)), ContainsNewlines::No).fmt(f)?;
+            }
+            Cow::Owned(normalized) => {
+                dynamic_text(&normalized, Some(range.start())).fmt(f)?;
+            }
+        }
+
+        text("j").fmt(f)
+    }
+}
+
+/// Returns the normalized integer string.
+fn normalize_integer(input: &str) -> Cow<str> {
+    // The normalized string if `input` is not yet normalized.
+    // `output` must remain empty if `input` is already normalized.
+    let mut output = String::new();
+    // Tracks the last index of `input` that has been written to `output`.
+    // If `last_index` is `0` at the end, then the input is already normalized and can be returned as is.
+    let mut last_index = 0;
+
+    let mut is_hex = false;
+
+    let mut chars = input.char_indices();
+
+    if let Some((_, '0')) = chars.next() {
+        if let Some((index, c)) = chars.next() {
+            is_hex = matches!(c, 'x' | 'X');
+            if matches!(c, 'B' | 'O' | 'X') {
+                // Lowercase the prefix.
+                output.push('0');
+                output.push(c.to_ascii_lowercase());
+                last_index = index + c.len_utf8();
+            }
+        }
+    }
+
+    // Skip the rest if `input` is not a hexinteger because there are only digits.
+    if is_hex {
+        for (index, c) in chars {
+            if matches!(c, 'a'..='f') {
+                // Uppercase hexdigits.
+                output.push_str(&input[last_index..index]);
+                output.push(c.to_ascii_uppercase());
+                last_index = index + c.len_utf8();
+            }
+        }
+    }
+
+    if last_index == 0 {
+        Cow::Borrowed(input)
+    } else {
+        output.push_str(&input[last_index..]);
+        Cow::Owned(output)
+    }
+}
+
+/// Returns the normalized floating number string.
+fn normalize_floating_number(input: &str) -> Cow<str> {
+    // The normalized string if `input` is not yet normalized.
+    // `output` must remain empty if `input` is already normalized.
+    let mut output = String::new();
+    // Tracks the last index of `input` that has been written to `output`.
+    // If `last_index` is `0` at the end, then the input is already normalized and can be returned as is.
+    let mut last_index = 0;
+
+    let mut chars = input.char_indices();
+
+    let fraction_ends_with_dot = if let Some((index, '.')) = chars.next() {
+        // Add a leading `0` if `input` starts with `.`.
+        output.push('0');
+        output.push('.');
+        last_index = index + '.'.len_utf8();
+        true
+    } else {
+        false
+    };
+
+    loop {
+        match chars.next() {
+            Some((index, c @ ('e' | 'E'))) => {
+                if fraction_ends_with_dot {
+                    // Add `0` if fraction part ends with `.`.
+                    output.push_str(&input[last_index..index]);
+                    output.push('0');
+                    last_index = index;
+                }
+
+                if c == 'E' {
+                    // Lowercase exponent part.
+                    output.push_str(&input[last_index..index]);
+                    output.push('e');
+                    last_index = index + 'E'.len_utf8();
+                }
+
+                if let Some((index, '+')) = chars.next() {
+                    // Remove `+` in exponent part.
+                    output.push_str(&input[last_index..index]);
+                    last_index = index + '+'.len_utf8();
+                }
+
+                break;
+            }
+            Some(_) => continue,
+            None => {
+                if input.ends_with('.') {
+                    // Add `0` if fraction part ends with `.`.
+                    output.push_str(&input[last_index..]);
+                    output.push('0');
+                    last_index = input.len();
+                }
+
+                break;
+            }
+        }
+    }
+
+    if last_index == 0 {
+        Cow::Borrowed(input)
+    } else {
+        output.push_str(&input[last_index..]);
+        Cow::Owned(output)
+    }
+}

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_36__numeric_literals.py.snap

@@ -1,118 +0,0 @@
----
-source: crates/ruff_python_formatter/tests/fixtures.rs
-input_file: crates/ruff_python_formatter/resources/test/fixtures/black/py_36/numeric_literals.py
----
-## Input
-
-```py
-#!/usr/bin/env python3.6
-
-x = 123456789
-x = 123456
-x = .1
-x = 1.
-x = 1E+1
-x = 1E-1
-x = 1.000_000_01
-x = 123456789.123456789
-x = 123456789.123456789E123456789
-x = 123456789E123456789
-x = 123456789J
-x = 123456789.123456789J
-x = 0XB1ACC
-x = 0B1011
-x = 0O777
-x = 0.000000006
-x = 10000
-x = 133333
-```
-
-## Black Differences
-
-```diff
---- Black
-+++ Ruff
-@@ -2,19 +2,19 @@
- 
- x = 123456789
- x = 123456
--x = 0.1
--x = 1.0
--x = 1e1
--x = 1e-1
-+x = .1
-+x = 1.
-+x = 1E+1
-+x = 1E-1
- x = 1.000_000_01
- x = 123456789.123456789
--x = 123456789.123456789e123456789
--x = 123456789e123456789
--x = 123456789j
--x = 123456789.123456789j
--x = 0xB1ACC
--x = 0b1011
--x = 0o777
-+x = 123456789.123456789E123456789
-+x = 123456789E123456789
-+x = 123456789J
-+x = 123456789.123456789J
-+x = 0XB1ACC
-+x = 0B1011
-+x = 0O777
- x = 0.000000006
- x = 10000
- x = 133333
-```
-
-## Ruff Output
-
-```py
-#!/usr/bin/env python3.6
-
-x = 123456789
-x = 123456
-x = .1
-x = 1.
-x = 1E+1
-x = 1E-1
-x = 1.000_000_01
-x = 123456789.123456789
-x = 123456789.123456789E123456789
-x = 123456789E123456789
-x = 123456789J
-x = 123456789.123456789J
-x = 0XB1ACC
-x = 0B1011
-x = 0O777
-x = 0.000000006
-x = 10000
-x = 133333
-```
-
-## Black Output
-
-```py
-#!/usr/bin/env python3.6
-
-x = 123456789
-x = 123456
-x = 0.1
-x = 1.0
-x = 1e1
-x = 1e-1
-x = 1.000_000_01
-x = 123456789.123456789
-x = 123456789.123456789e123456789
-x = 123456789e123456789
-x = 123456789j
-x = 123456789.123456789j
-x = 0xB1ACC
-x = 0b1011
-x = 0o777
-x = 0.000000006
-x = 10000
-x = 133333
-```
-
-

crates/ruff_python_formatter/tests/snapshots/black_compatibility@py_36__numeric_literals_skip_underscores.py.snap

@@ -1,72 +0,0 @@
----
-source: crates/ruff_python_formatter/tests/fixtures.rs
-input_file: crates/ruff_python_formatter/resources/test/fixtures/black/py_36/numeric_literals_skip_underscores.py
----
-## Input
-
-```py
-#!/usr/bin/env python3.6
-
-x = 123456789
-x = 1_2_3_4_5_6_7
-x = 1E+1
-x = 0xb1acc
-x = 0.00_00_006
-x = 12_34_567J
-x = .1_2
-x = 1_2.
-```
-
-## Black Differences
-
-```diff
---- Black
-+++ Ruff
-@@ -2,9 +2,9 @@
- 
- x = 123456789
- x = 1_2_3_4_5_6_7
--x = 1e1
--x = 0xB1ACC
-+x = 1E+1
-+x = 0xb1acc
- x = 0.00_00_006
--x = 12_34_567j
--x = 0.1_2
--x = 1_2.0
-+x = 12_34_567J
-+x = .1_2
-+x = 1_2.
-```
-
-## Ruff Output
-
-```py
-#!/usr/bin/env python3.6
-
-x = 123456789
-x = 1_2_3_4_5_6_7
-x = 1E+1
-x = 0xb1acc
-x = 0.00_00_006
-x = 12_34_567J
-x = .1_2
-x = 1_2.
-```
-
-## Black Output
-
-```py
-#!/usr/bin/env python3.6
-
-x = 123456789
-x = 1_2_3_4_5_6_7
-x = 1e1
-x = 0xB1ACC
-x = 0.00_00_006
-x = 12_34_567j
-x = 0.1_2
-x = 1_2.0
-```
-
-

crates/ruff_python_formatter/tests/snapshots/black_compatibility@simple_cases__attribute_access_on_number_literals.py.snap

@@ -34,38 +34,21 @@ y = 100(no)
 ```diff
 --- Black
 +++ Ruff
-@@ -1,19 +1,19 @@
+@@ -8,10 +8,10 @@
- x = (123456789).bit_count()
+ x = (123456789.123456789e123456789).real
- x = (123456).__abs__()
+ x = (123456789e123456789).conjugate()
--x = (0.1).is_integer()
+ x = 123456789j.real
--x = (1.0).imag
--x = (1e1).imag
--x = (1e-1).real
-+x = (.1).is_integer()
-+x = (1.).imag
-+x = (1E+1).imag
-+x = (1E-1).real
- x = (123456789.123456789).hex()
--x = (123456789.123456789e123456789).real
--x = (123456789e123456789).conjugate()
--x = 123456789j.real
 -x = 123456789.123456789j.__add__(0b1011.bit_length())
 -x = 0xB1ACC.conjugate()
 -x = 0b1011.conjugate()
 -x = 0o777.real
-+x = (123456789.123456789E123456789).real
++x = 123456789.123456789j.__add__((0b1011).bit_length())
-+x = (123456789E123456789).conjugate()
++x = (0xB1ACC).conjugate()
-+x = 123456789J.real
++x = (0b1011).conjugate()
-+x = 123456789.123456789J.__add__((0b1011).bit_length())
++x = (0o777).real
-+x = (0XB1ACC).conjugate()
-+x = (0B1011).conjugate()
-+x = (0O777).real
  x = (0.000000006).hex()
--x = -100.0000j
+ x = -100.0000j
-+x = -100.0000J
  
- if (10).real:
-     ...
 ```
 
 ## Ruff Output
@@ -73,20 +56,20 @@ y = 100(no)
 ```py
 x = (123456789).bit_count()
 x = (123456).__abs__()
-x = (.1).is_integer()
+x = (0.1).is_integer()
-x = (1.).imag
+x = (1.0).imag
-x = (1E+1).imag
+x = (1e1).imag
-x = (1E-1).real
+x = (1e-1).real
 x = (123456789.123456789).hex()
-x = (123456789.123456789E123456789).real
+x = (123456789.123456789e123456789).real
-x = (123456789E123456789).conjugate()
+x = (123456789e123456789).conjugate()
-x = 123456789J.real
+x = 123456789j.real
-x = 123456789.123456789J.__add__((0b1011).bit_length())
+x = 123456789.123456789j.__add__((0b1011).bit_length())
-x = (0XB1ACC).conjugate()
+x = (0xB1ACC).conjugate()
-x = (0B1011).conjugate()
+x = (0b1011).conjugate()
-x = (0O777).real
+x = (0o777).real
 x = (0.000000006).hex()
-x = -100.0000J
+x = -100.0000j
 
 if (10).real:
     ...