Implement complex literal formatting (#3186)

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::{float_literal, int_literal};
+use crate::format::numbers::{complex_literal, float_literal, int_literal};
 use crate::format::strings::string_literal;
 use crate::shared_traits::AsFormat;
 use crate::trivia::{Parenthesize, Relationship, TriviaKind};
@@ -645,6 +645,7 @@ fn format_constant(
     _kind: Option<&str>,
 ) -> FormatResult<()> {
     match constant {
+        Constant::Ellipsis => write!(f, [text("...")])?,
         Constant::None => write!(f, [text("None")])?,
         Constant::Bool(value) => {
             if *value {
@@ -655,8 +656,10 @@ 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))])?,
+        Constant::Str(_) => write!(f, [string_literal(expr)])?,
+        Constant::Bytes(_) => write!(f, [string_literal(expr)])?,
+        Constant::Complex { .. } => write!(f, [complex_literal(Range::from_located(expr))])?,
+        Constant::Tuple(_) => unreachable!("Constant::Tuple should be handled by format_tuple"),
     }
     Ok(())
 }

crates/ruff_python_formatter/src/format/numbers.rs

@@ -162,3 +162,39 @@ impl Format<ASTFormatContext<'_>> for IntLiteral {
 pub const fn int_literal(range: Range) -> IntLiteral {
     IntLiteral { range }
 }
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub struct ComplexLiteral {
+    range: Range,
+}
+
+impl Format<ASTFormatContext<'_>> for ComplexLiteral {
+    fn fmt(&self, f: &mut Formatter<ASTFormatContext<'_>>) -> FormatResult<()> {
+        let (source, start, end) = f.context().locator().slice(self.range);
+
+        if source[start..end].ends_with('j') {
+            write!(f, [literal(self.range)])?;
+        } else if source[start..end].ends_with('J') {
+            write!(
+                f,
+                [literal(Range::new(
+                    self.range.location,
+                    Location::new(
+                        self.range.end_location.row(),
+                        self.range.end_location.column() - 1
+                    ),
+                ))]
+            )?;
+            write!(f, [text("j")])?;
+        } else {
+            unreachable!("expected complex literal to end with j or J");
+        }
+
+        Ok(())
+    }
+}
+
+#[inline]
+pub const fn complex_literal(range: Range) -> ComplexLiteral {
+    ComplexLiteral { range }
+}

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__attribute_access_on_number_literals_py.snap

@@ -46,8 +46,6 @@ y = 100(no)
 -x = (123456789.123456789).hex()
 -x = (123456789.123456789e123456789).real
 -x = (123456789e123456789).conjugate()
--x = 123456789j.real
--x = 123456789.123456789j.__add__(0b1011.bit_length())
 +x = 0.1.is_integer()
 +x = 1.0.imag
 +x = 1e1.imag
@@ -55,15 +53,14 @@ y = 100(no)
 +x = 123456789.123456789.hex()
 +x = 123456789.123456789e123456789.real
 +x = 123456789e123456789.conjugate()
-+x = 123456789J.real
-+x = 123456789.123456789J.__add__(0b1011.bit_length())
+ x = 123456789j.real
+ x = 123456789.123456789j.__add__(0b1011.bit_length())
  x = 0xB1ACC.conjugate()
  x = 0b1011.conjugate()
  x = 0o777.real
 -x = (0.000000006).hex()
--x = -100.0000j
 +x = 0.000000006.hex()
-+x = -100.0000J
+ x = -100.0000j
  
 -if (10).real:
 +if 10.real:
@@ -86,13 +83,13 @@ 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 = 123456789j.real
+x = 123456789.123456789j.__add__(0b1011.bit_length())
 x = 0xB1ACC.conjugate()
 x = 0b1011.conjugate()
 x = 0o777.real
 x = 0.000000006.hex()
-x = -100.0000J
+x = -100.0000j
 
 if 10.real:
     ...