Format ExprIfExp (ternary operator)

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/if.py

@@ -0,0 +1,33 @@
+a1 = 1 if True else 2
+
+a2 = "this is a very long text that will make the group break to check that parentheses are added" if True else 2
+
+# These comment should be kept in place
+b1 = (
+ # We return "a" ...
+ "a" # that's our True value
+ # ... if this condition matches ...
+ if True # that's our test
+ # ... otherwise we return "b"
+ else "b" # that's our False value
+)
+
+# These only need to be stable, bonus is we also keep the order
+c1 = (
+ "a" # 1
+ if # 2
+ True # 3
+ else # 4
+ "b" # 5
+)
+c2 = (
+ "a" # 1
+ # 2
+ if # 3
+ # 4
+ True # 5
+ # 6
+ else # 7
+ # 8
+ "b" # 9
+)

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,8 +1,8 @@
 use std::cmp::Ordering;
 
-use ruff_text_size::TextRange;
+use ruff_text_size::{TextLen, TextRange};
 use rustpython_parser::ast;
-use rustpython_parser::ast::{Expr, ExprSlice, Ranged};
+use rustpython_parser::ast::{Expr, ExprIfExp, ExprSlice, Ranged};
 
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
 use ruff_python_ast::source_code::Locator;
@@ -14,7 +14,9 @@ use crate::expression::expr_slice::{assign_comment_in_slice, ExprSliceCommentSec
 use crate::other::arguments::{
  assign_argument_separator_comment_placement, find_argument_separators,
 };
-use crate::trivia::{first_non_trivia_token_rev, SimpleTokenizer, Token, TokenKind};
+use crate::trivia::{
+ first_non_trivia_token, first_non_trivia_token_rev, SimpleTokenizer, Token, TokenKind,
+};
 
 /// Implements the custom comment placement logic.
 pub(super) fn place_comment<'a>(
@@ -37,6 +39,7 @@ pub(super) fn place_comment<'a>(
  handle_dict_unpacking_comment,
  handle_slice_comments,
  handle_attribute_comment,
+ handle_expr_if_comment,
  ];
  for handler in HANDLERS {
  comment = match handler(comment, locator) {
@@ -1154,6 +1157,80 @@ fn handle_attribute_comment<'a>(
  }
 }
 
+/// Assign comments between `if` and `test` and `else` and `orelse` as leading to the respective
+/// node.
+///
+/// ```python
+/// x = (
+/// "a"
+/// if # leading comment of `True`
+/// True
+/// else # leading comment of `"b"`
+/// "b"
+/// )
+/// ```
+///
+/// This placement ensures comments remain in their previous order. This an edge case that only
+/// happens if the comments are in a weird position but it also doesn't hurt handling it.
+fn handle_expr_if_comment<'a>(
+ comment: DecoratedComment<'a>,
+ locator: &Locator,
+) -> CommentPlacement<'a> {
+ let Some(expr_if) = comment.enclosing_node().expr_if_exp() else {
+ return CommentPlacement::Default(comment);
+ };
+ let ExprIfExp {
+ range: _,
+ test,
+ body,
+ orelse,
+ } = expr_if;
+
+ // Find the if and the else
+ let if_token =
+ find_only_token_str_in_range(TextRange::new(body.end(), test.start()), locator, "if");
+ let else_token =
+ find_only_token_str_in_range(TextRange::new(test.end(), orelse.start()), locator, "else");
+
+ if if_token.range.start() < comment.slice().start()
+ && comment.slice().start() < test.start()
+ && comment.line_position().is_end_of_line()
+ {
+ return CommentPlacement::leading(test.as_ref().into(), comment);
+ }
+
+ if else_token.range.start() < comment.slice().start()
+ && comment.slice().start() < orelse.start()
+ && comment.line_position().is_end_of_line()
+ {
+ return CommentPlacement::leading(orelse.as_ref().into(), comment);
+ }
+
+ CommentPlacement::Default(comment)
+}
+
+/// Looks for a multi char token in the range that contains no other tokens. `SimpleTokenizer` only
+/// works with single char tokens so we check that we have the right token by string comparison.
+fn find_only_token_str_in_range(range: TextRange, locator: &Locator, token_str: &str) -> Token {
+ let token =
+ first_non_trivia_token(range.start(), locator.contents()).expect("Expected a token");
+ debug_assert!(
+ locator.after(token.start()).starts_with(token_str),
+ "expected a `{token_str}` token",
+ );
+ debug_assert!(
+ SimpleTokenizer::new(
+ locator.contents(),
+ TextRange::new(token.start() + token_str.text_len(), range.end())
+ )
+ .skip_trivia()
+ .next()
+ .is_none(),
+ "Didn't expect any other token"
+ );
+ token
+}
+
 /// Returns `true` if `right` is `Some` and `left` and `right` are referentially equal.
 fn are_same_optional<'a, T>(left: AnyNodeRef, right: Option<T>) -> bool
 where

crates/ruff_python_formatter/src/expression/expr_if_exp.rs

@@ -1,21 +1,42 @@
-use crate::comments::Comments;
+use crate::comments::{leading_comments, Comments};
 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 crate::{AsFormat, FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::{group, soft_line_break_or_space, space, text};
+use ruff_formatter::{format_args, write, Buffer, FormatResult};
 use rustpython_parser::ast::ExprIfExp;
 
 #[derive(Default)]
 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<()> {
+ let ExprIfExp {
+ range: _,
+ test,
+ body,
+ orelse,
+ } = item;
+ let comments = f.context().comments().clone();
+ // We place `if test` and `else orelse` on a single line, so the `test` and `orelse` leading
+ // comments go on the line before the `if` or `else` instead of directly ahead `test` or
+ // `orelse`
  write!(
  f,
- [not_yet_implemented_custom_text(
- "NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false"
- )]
+ [group(&format_args![
+ body.format(),
+ soft_line_break_or_space(),
+ leading_comments(comments.leading_comments(test.as_ref())),
+ text("if"),
+ space(),
+ test.format(),
+ soft_line_break_or_space(),
+ leading_comments(comments.leading_comments(orelse.as_ref())),
+ text("else"),
+ space(),
+ orelse.format()
+ ])]
  )
  }
 }

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

@@ -79,15 +79,17 @@ def something():
 ```diff
 --- Black
 +++ Ruff
-@@ -1,90 +1,50 @@
+@@ -1,20 +1,16 @@
  long_kwargs_single_line = my_function(
  foo="test, this is a sample value",
 - bar=(
 - some_long_value_name_foo_bar_baz
 - if some_boolean_variable
 - else some_fallback_value_foo_bar_baz
 - ),
-+ bar=NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false,
++ bar=some_long_value_name_foo_bar_baz
++ if some_boolean_variable
++ else some_fallback_value_foo_bar_baz,
  baz="hello, this is a another value",
  )
  
@@ -98,33 +100,13 @@ def something():
 - if some_boolean_variable
 - else some_fallback_value_foo_bar_baz
 - ),
-+ bar=NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false,
++ bar=some_long_value_name_foo_bar_baz
++ if some_boolean_variable
++ else some_fallback_value_foo_bar_baz,
  baz="hello, this is a another value",
  )
  
- imploding_kwargs = my_function(
- foo="test, this is a sample value",
-- bar=a if foo else b,
-+ bar=NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false,
- baz="hello, this is a another value",
- )
- 
--imploding_line = 1 if 1 + 1 == 2 else 0
-+imploding_line = NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
- 
--exploding_line = (
-- "hello this is a slightly long string"
-- if some_long_value_name_foo_bar_baz
-- else "this one is a little shorter"
--)
-+exploding_line = NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
- 
- positional_argument_test(
-- some_long_value_name_foo_bar_baz
-- if some_boolean_variable
-- else some_fallback_value_foo_bar_baz
-+ NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
- )
+@@ -40,11 +36,9 @@
  
  
  def weird_default_argument(
@@ -133,21 +115,15 @@ def something():
 - if SOME_CONSTANT
 - else some_fallback_value_foo_bar_baz
 - ),
-+ x=NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
++ x=some_long_value_name_foo_bar_baz
++ if SOME_CONSTANT
++ else some_fallback_value_foo_bar_baz,
  ):
  pass
  
- 
--nested = (
-- "hello this is a slightly long string"
-- if (
-- some_long_value_name_foo_bar_baz
-- if nesting_test_expressions
-- else some_fallback_value_foo_bar_baz
-- )
-- else "this one is a little shorter"
--)
-+nested = NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
+@@ -59,26 +53,14 @@
+ else "this one is a little shorter"
+ )
  
 -generator_expression = (
 - (
@@ -174,52 +150,65 @@ def something():
  )
  
  
- def something():
- clone._iterable_class = (
-- NamedValuesListIterable
-- if named
-- else FlatValuesListIterable if flat else ValuesListIterable
-+ NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
- )
 ```
 
 ## Ruff Output
 
 ```py
 long_kwargs_single_line = my_function(
  foo="test, this is a sample value",
- bar=NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false,
+ bar=some_long_value_name_foo_bar_baz
+ if some_boolean_variable
+ else some_fallback_value_foo_bar_baz,
  baz="hello, this is a another value",
 )
 
 multiline_kwargs_indented = my_function(
  foo="test, this is a sample value",
- bar=NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false,
+ bar=some_long_value_name_foo_bar_baz
+ if some_boolean_variable
+ else some_fallback_value_foo_bar_baz,
  baz="hello, this is a another value",
 )
 
 imploding_kwargs = my_function(
  foo="test, this is a sample value",
- bar=NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false,
+ bar=a if foo else b,
  baz="hello, this is a another value",
 )
 
-imploding_line = NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
+imploding_line = 1 if 1 + 1 == 2 else 0
 
-exploding_line = NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
+exploding_line = (
+ "hello this is a slightly long string"
+ if some_long_value_name_foo_bar_baz
+ else "this one is a little shorter"
+)
 
 positional_argument_test(
- NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
+ some_long_value_name_foo_bar_baz
+ if some_boolean_variable
+ else some_fallback_value_foo_bar_baz
 )
 
 
 def weird_default_argument(
- x=NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
+ x=some_long_value_name_foo_bar_baz
+ if SOME_CONSTANT
+ else some_fallback_value_foo_bar_baz,
 ):
  pass
 
 
-nested = NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
+nested = (
+ "hello this is a slightly long string"
+ if (
+ some_long_value_name_foo_bar_baz
+ if nesting_test_expressions
+ else some_fallback_value_foo_bar_baz
+ )
+ else "this one is a little shorter"
+)
 
 generator_expression = (NOT_YET_IMPLEMENTED_generator_key for NOT_YET_IMPLEMENTED_generator_key in [])
 
@@ -234,7 +223,9 @@ def limit_offset_sql(self, low_mark, high_mark):
 
 def something():
  clone._iterable_class = (
- NOT_IMPLEMENTED_true if NOT_IMPLEMENTED_cond else NOT_IMPLEMENTED_false
+ NamedValuesListIterable
+ if named
+ else FlatValuesListIterable if flat else ValuesListIterable
  )
 ```