Format numeric constants (#5972)

Co-authored-by: Micha Reiser <micha@reiser.io>
This commit is contained in:
Luc Khai Hai 2023-07-24 16:04:40 +09:00 committed by GitHub
parent 33196f1859
commit dfa81b6fe0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 227 additions and 233 deletions

View file

@ -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 { .. } => {
write!(f, [verbatim_text(item)])
}
Constant::Int(_) => FormatInt::new(item).fmt(f),
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)

View file

@ -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;

View file

@ -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)
}
}