Add basic Constant formatting (#4954)

astral-sh · Jun 8, 2023 · c1cc6f3 · c1cc6f3
1 parent 83cf6d6
commit c1cc6f3
Show file tree

Hide file tree

Showing 47 changed files with 548 additions and 623 deletions.
diff --git a/crates/ruff_python_formatter/src/builders.rs b/crates/ruff_python_formatter/src/builders.rs
@@ -191,16 +191,16 @@ no_leading_newline = 30
 
         assert_eq!(
             &printed,
-            r#"a = 0x42
+            r#"a = 10
 
 
-three_leading_newlines = 0x42
+three_leading_newlines = 80
 
 
-two_leading_newlines = 0x42
+two_leading_newlines = 20
 
-one_leading_newline = 0x42
-no_leading_newline = 0x42"#
+one_leading_newline = 10
+no_leading_newline = 30"#
         );
     }
 
@@ -211,14 +211,14 @@ no_leading_newline = 0x42"#
 
         assert_eq!(
             &printed,
-            r#"a = 0x42
+            r#"a = 10
 
-three_leading_newlines = 0x42
+three_leading_newlines = 80
 
-two_leading_newlines = 0x42
+two_leading_newlines = 20
 
-one_leading_newline = 0x42
-no_leading_newline = 0x42"#
+one_leading_newline = 10
+no_leading_newline = 30"#
         );
     }
 
@@ -229,11 +229,11 @@ no_leading_newline = 0x42"#
 
         assert_eq!(
             &printed,
-            r#"a = 0x42
-three_leading_newlines = 0x42
-two_leading_newlines = 0x42
-one_leading_newline = 0x42
-no_leading_newline = 0x42"#
+            r#"a = 10
+three_leading_newlines = 80
+two_leading_newlines = 20
+one_leading_newline = 10
+no_leading_newline = 30"#
         );
     }
 }
diff --git a/crates/ruff_python_formatter/src/expression/expr_attribute.rs b/crates/ruff_python_formatter/src/expression/expr_attribute.rs
@@ -1,24 +1,38 @@
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
-use crate::{not_yet_implemented_custom_text, AsFormat, FormatNodeRule, PyFormatter};
-use ruff_formatter::prelude::text;
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::ExprAttribute;
+use crate::prelude::*;
+use crate::{not_yet_implemented_custom_text, FormatNodeRule};
+use ruff_formatter::write;
+use rustpython_parser::ast::{Constant, Expr, ExprAttribute, ExprConstant};
 
 #[derive(Default)]
 pub struct FormatExprAttribute;
 
 impl FormatNodeRule<ExprAttribute> for FormatExprAttribute {
     fn fmt_fields(&self, item: &ExprAttribute, f: &mut PyFormatter) -> FormatResult<()> {
-        // We need to write the value - which is also a dummy - already because power op spacing
-        // depends on it
+        let ExprAttribute {
+            value,
+            range: _,
+            attr: _,
+            ctx: _,
+        } = item;
+
+        let requires_space = matches!(
+            value.as_ref(),
+            Expr::Constant(ExprConstant {
+                value: Constant::Int(_) | Constant::Float(_),
+                ..
+            })
+        );
+
         write!(
             f,
             [
                 item.value.format(),
+                requires_space.then_some(space()),
                 text("."),
-                not_yet_implemented_custom_text(item, "NOT_IMPLEMENTED_attr")
+                not_yet_implemented_custom_text("NOT_IMPLEMENTED_attr")
             ]
         )
     }

diff --git a/crates/ruff_python_formatter/src/expression/expr_bool_op.rs b/crates/ruff_python_formatter/src/expression/expr_bool_op.rs
@@ -9,11 +9,10 @@ use rustpython_parser::ast::ExprBoolOp;
 pub struct FormatExprBoolOp;
 
 impl FormatNodeRule<ExprBoolOp> for FormatExprBoolOp {
-    fn fmt_fields(&self, item: &ExprBoolOp, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, _item: &ExprBoolOp, f: &mut PyFormatter) -> FormatResult<()> {
         write!(
             f,
             [not_yet_implemented_custom_text(
-                item,
                 "NOT_IMPLEMENTED_bool_op1 and NOT_IMPLEMENTED_bool_op2"
             )]
         )

diff --git a/crates/ruff_python_formatter/src/expression/expr_call.rs b/crates/ruff_python_formatter/src/expression/expr_call.rs
@@ -9,13 +9,10 @@ use rustpython_parser::ast::ExprCall;
 pub struct FormatExprCall;
 
 impl FormatNodeRule<ExprCall> for FormatExprCall {
-    fn fmt_fields(&self, item: &ExprCall, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, _item: &ExprCall, f: &mut PyFormatter) -> FormatResult<()> {
         write!(
             f,
-            [not_yet_implemented_custom_text(
-                item,
-                "NOT_IMPLEMENTED_call()"
-            )]
+            [not_yet_implemented_custom_text("NOT_IMPLEMENTED_call()")]
         )
     }
 }

diff --git a/crates/ruff_python_formatter/src/expression/expr_compare.rs b/crates/ruff_python_formatter/src/expression/expr_compare.rs
@@ -10,11 +10,10 @@ use rustpython_parser::ast::ExprCompare;
 pub struct FormatExprCompare;
 
 impl FormatNodeRule<ExprCompare> for FormatExprCompare {
-    fn fmt_fields(&self, item: &ExprCompare, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, _item: &ExprCompare, f: &mut PyFormatter) -> FormatResult<()> {
         write!(
             f,
             [not_yet_implemented_custom_text(
-                item,
                 "NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right"
             )]
         )

diff --git a/crates/ruff_python_formatter/src/expression/expr_constant.rs b/crates/ruff_python_formatter/src/expression/expr_constant.rs
@@ -1,16 +1,42 @@
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
-use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::ExprConstant;
+use crate::prelude::*;
+use crate::{not_yet_implemented_custom_text, verbatim_text, FormatNodeRule};
+use ruff_formatter::write;
+use rustpython_parser::ast::{Constant, ExprConstant};
 
 #[derive(Default)]
 pub struct FormatExprConstant;
 
 impl FormatNodeRule<ExprConstant> for FormatExprConstant {
     fn fmt_fields(&self, item: &ExprConstant, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented_custom_text(item, "0x42")])
+        let ExprConstant {
+            range: _,
+            value,
+            kind: _,
+        } = item;
+
+        match value {
+            Constant::Ellipsis => text("...").fmt(f),
+            Constant::None => text("None").fmt(f),
+            Constant::Bool(value) => match value {
+                true => text("True").fmt(f),
+                false => text("False").fmt(f),
+            },
+            Constant::Int(_) | Constant::Float(_) | Constant::Complex { .. } => {
+                write!(f, [verbatim_text(item)])
+            }
+            Constant::Str(_) => {
+                not_yet_implemented_custom_text(r#""NOT_YET_IMPLEMENTED_STRING""#).fmt(f)
+            }
+            Constant::Bytes(_) => {
+                not_yet_implemented_custom_text(r#"b"NOT_YET_IMPLEMENTED_BYTE_STRING""#).fmt(f)
+            }
+            Constant::Tuple(_) => {
+                not_yet_implemented_custom_text("(NOT_YET_IMPLEMENTED_TUPLE,)").fmt(f)
+            }
+        }
     }
 }
 

diff --git a/crates/ruff_python_formatter/src/expression/expr_dict.rs b/crates/ruff_python_formatter/src/expression/expr_dict.rs
@@ -9,11 +9,10 @@ use rustpython_parser::ast::ExprDict;
 pub struct FormatExprDict;
 
 impl FormatNodeRule<ExprDict> for FormatExprDict {
-    fn fmt_fields(&self, item: &ExprDict, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, _item: &ExprDict, f: &mut PyFormatter) -> FormatResult<()> {
         write!(
             f,
             [not_yet_implemented_custom_text(
-                item,
                 "{NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value}"
             )]
         )

diff --git a/crates/ruff_python_formatter/src/expression/expr_dict_comp.rs b/crates/ruff_python_formatter/src/expression/expr_dict_comp.rs
@@ -9,11 +9,10 @@ use rustpython_parser::ast::ExprDictComp;
 pub struct FormatExprDictComp;
 
 impl FormatNodeRule<ExprDictComp> for FormatExprDictComp {
-    fn fmt_fields(&self, item: &ExprDictComp, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, _item: &ExprDictComp, f: &mut PyFormatter) -> FormatResult<()> {
         write!(
             f,
             [not_yet_implemented_custom_text(
-                item,
                 "{NOT_IMPLEMENTED_dict_key: NOT_IMPLEMENTED_dict_value for key, value in NOT_IMPLEMENTED_dict}"
             )]
         )

diff --git a/crates/ruff_python_formatter/src/expression/expr_generator_exp.rs b/crates/ruff_python_formatter/src/expression/expr_generator_exp.rs
@@ -9,11 +9,8 @@ use rustpython_parser::ast::ExprGeneratorExp;
 pub struct FormatExprGeneratorExp;
 
 impl FormatNodeRule<ExprGeneratorExp> for FormatExprGeneratorExp {
-    fn fmt_fields(&self, item: &ExprGeneratorExp, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(
-            f,
-            [not_yet_implemented_custom_text(item, "(i for i in [])")]
-        )
+    fn fmt_fields(&self, _item: &ExprGeneratorExp, f: &mut PyFormatter) -> FormatResult<()> {
+        write!(f, [not_yet_implemented_custom_text("(i for i in [])")])
     }
 }
 

diff --git a/crates/ruff_python_formatter/src/expression/expr_if_exp.rs b/crates/ruff_python_formatter/src/expression/expr_if_exp.rs
@@ -9,11 +9,10 @@ use rustpython_parser::ast::ExprIfExp;
 pub struct FormatExprIfExp;
 
 impl FormatNodeRule<ExprIfExp> for FormatExprIfExp {
-    fn fmt_fields(&self, item: &ExprIfExp, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, _item: &ExprIfExp, f: &mut PyFormatter) -> FormatResult<()> {
         write!(
             f,
             [not_yet_implemented_custom_text(
-                item,
                 "NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false"
             )]
         )

diff --git a/crates/ruff_python_formatter/src/expression/expr_lambda.rs b/crates/ruff_python_formatter/src/expression/expr_lambda.rs
@@ -9,8 +9,8 @@ use rustpython_parser::ast::ExprLambda;
 pub struct FormatExprLambda;
 
 impl FormatNodeRule<ExprLambda> for FormatExprLambda {
-    fn fmt_fields(&self, item: &ExprLambda, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented_custom_text(item, "lambda x: True")])
+    fn fmt_fields(&self, _item: &ExprLambda, f: &mut PyFormatter) -> FormatResult<()> {
+        write!(f, [not_yet_implemented_custom_text("lambda x: True")])
     }
 }
 

diff --git a/crates/ruff_python_formatter/src/expression/expr_list_comp.rs b/crates/ruff_python_formatter/src/expression/expr_list_comp.rs
@@ -9,11 +9,8 @@ use rustpython_parser::ast::ExprListComp;
 pub struct FormatExprListComp;
 
 impl FormatNodeRule<ExprListComp> for FormatExprListComp {
-    fn fmt_fields(&self, item: &ExprListComp, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(
-            f,
-            [not_yet_implemented_custom_text(item, "[i for i in []]")]
-        )
+    fn fmt_fields(&self, _item: &ExprListComp, f: &mut PyFormatter) -> FormatResult<()> {
+        write!(f, [not_yet_implemented_custom_text("[i for i in []]")])
     }
 }
 

diff --git a/crates/ruff_python_formatter/src/expression/expr_slice.rs b/crates/ruff_python_formatter/src/expression/expr_slice.rs
@@ -10,11 +10,10 @@ use rustpython_parser::ast::ExprSlice;
 pub struct FormatExprSlice;
 
 impl FormatNodeRule<ExprSlice> for FormatExprSlice {
-    fn fmt_fields(&self, item: &ExprSlice, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, _item: &ExprSlice, f: &mut PyFormatter) -> FormatResult<()> {
         write!(
             f,
             [not_yet_implemented_custom_text(
-                item,
                 "NOT_IMPLEMENTED_start:NOT_IMPLEMENTED_end"
             )]
         )

diff --git a/crates/ruff_python_formatter/src/expression/expr_subscript.rs b/crates/ruff_python_formatter/src/expression/expr_subscript.rs
@@ -10,11 +10,10 @@ use rustpython_parser::ast::ExprSubscript;
 pub struct FormatExprSubscript;
 
 impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
-    fn fmt_fields(&self, item: &ExprSubscript, f: &mut PyFormatter) -> FormatResult<()> {
+    fn fmt_fields(&self, _item: &ExprSubscript, f: &mut PyFormatter) -> FormatResult<()> {
         write!(
             f,
             [not_yet_implemented_custom_text(
-                item,
                 "NOT_IMPLEMENTED_value[NOT_IMPLEMENTED_key]"
             )]
         )

diff --git a/crates/ruff_python_formatter/src/expression/expr_tuple.rs b/crates/ruff_python_formatter/src/expression/expr_tuple.rs
@@ -9,8 +9,8 @@ use rustpython_parser::ast::ExprTuple;
 pub struct FormatExprTuple;
 
 impl FormatNodeRule<ExprTuple> for FormatExprTuple {
-    fn fmt_fields(&self, item: &ExprTuple, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented_custom_text(item, "(1, 2)")])
+    fn fmt_fields(&self, _item: &ExprTuple, f: &mut PyFormatter) -> FormatResult<()> {
+        write!(f, [not_yet_implemented_custom_text("(1, 2)")])
     }
 }
 

diff --git a/crates/ruff_python_formatter/src/lib.rs b/crates/ruff_python_formatter/src/lib.rs
@@ -39,19 +39,17 @@ where
     N: AstNode,
 {
     fn fmt(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [source_position(node.start())])?;
         self.fmt_leading_comments(node, f)?;
         self.fmt_node(node, f)?;
         self.fmt_dangling_comments(node, f)?;
-        self.fmt_trailing_comments(node, f)?;
-        write!(f, [source_position(node.end())])
+        self.fmt_trailing_comments(node, f)
     }
 
     /// Formats the node without comments. Ignores any suppression comments.
     fn fmt_node(&self, node: &N, f: &mut PyFormatter) -> FormatResult<()> {
+        write!(f, [source_position(node.start())])?;
         self.fmt_fields(node, f)?;
-
-        Ok(())
+        write!(f, [source_position(node.end())])
     }
 
     /// Formats the node's fields.
@@ -164,33 +162,26 @@ impl Format<PyFormatContext<'_>> for NotYetImplemented {
     }
 }
 
-pub(crate) struct NotYetImplementedCustomText(NodeKind, String);
+pub(crate) struct NotYetImplementedCustomText(&'static str);
 
 /// Formats a placeholder for nodes that have not yet been implemented
-pub(crate) fn not_yet_implemented_custom_text<'a, T>(
-    node: T,
-    text: impl AsRef<str>,
-) -> NotYetImplementedCustomText
-where
-    T: Into<AnyNodeRef<'a>>,
-{
-    NotYetImplementedCustomText(node.into().kind(), text.as_ref().to_string())
+pub(crate) const fn not_yet_implemented_custom_text(
+    text: &'static str,
+) -> NotYetImplementedCustomText {
+    NotYetImplementedCustomText(text)
 }
 
 impl Format<PyFormatContext<'_>> for NotYetImplementedCustomText {
     fn fmt(&self, f: &mut PyFormatter) -> FormatResult<()> {
         f.write_element(FormatElement::Tag(Tag::StartVerbatim(
             tag::VerbatimKind::Verbatim {
-                length: self.1.text_len(),
+                length: self.0.text_len(),
             },
         )))?;
 
-        f.write_element(FormatElement::DynamicText {
-            text: Box::from(self.1.clone()),
-        })?;
+        text(self.0).fmt(f)?;
 
-        f.write_element(FormatElement::Tag(Tag::EndVerbatim))?;
-        Ok(())
+        f.write_element(FormatElement::Tag(Tag::EndVerbatim))
     }
 }
 
@@ -417,12 +408,11 @@ Formatted twice:
     #[ignore]
     #[test]
     fn quick_test() {
-        let src = r#"AAAAAAAAAAAAA = AAAAAAAAAAAAA  # type: ignore
+        let src = r#"
+def test(): ...
 
-call_to_some_function_asdf(
-    foo,
-    [AAAAAAAAAAAAAAAAAAAAAAA, AAAAAAAAAAAAAAAAAAAAAAA, AAAAAAAAAAAAAAAAAAAAAAA, BBBBBBBBBBBB],  # type: ignore
-)
+# Comment
+def with_leading_comment(): ...
 "#;
         // Tokenize once
         let mut tokens = Vec::new();