Format Compare Op

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

@@ -0,0 +1,61 @@
+a == b
+a != b
+a < b
+a <= b
+a > b
+a >= b
+a is b
+a is not b
+a in b
+a not in b
+
+(a ==
+  # comment
+    b
+)
+
+(a == # comment
+ b
+ )
+
+a < b > c == d
+
+aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa < bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb > ccccccccccccccccccccccccccccc == ddddddddddddddddddddd
+
+aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa < [
+    bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb,
+    ff,
+] < [ccccccccccccccccccccccccccccc, dddd] < ddddddddddddddddddddddddddddddddddddddddddd
+
+return 1 == 2 and (
+    name,
+    description,
+    self_default,
+    self_selected,
+    self_auto_generated,
+    self_parameters,
+    self_meta_data,
+    self_schedule,
+) == (
+    name,
+    description,
+    othr_default,
+    othr_selected,
+    othr_auto_generated,
+    othr_parameters,
+    othr_meta_data,
+    othr_schedule,
+)
+
+(name, description, self_default, self_selected, self_auto_generated, self_parameters, self_meta_data, self_schedule) == (name, description, other_default, othr_selected, othr_auto_generated, othr_parameters, othr_meta_data, othr_schedule)
+((name, description, self_default, self_selected, self_auto_generated, self_parameters, self_meta_data, self_schedule) == (name, description, other_default, othr_selected, othr_auto_generated, othr_parameters, othr_meta_data, othr_schedule))
+
+[
+    (
+        a
+        + [
+            bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb
+        ]
+        >= c
+    )
+]

crates/ruff_python_formatter/src/expression/expr_compare.rs

@@ -1,23 +1,96 @@
+use crate::comments::{leading_comments, Comments};
+use crate::expression::binary_like::{BinaryLayout, FormatBinaryLike};
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
-use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
-
-use crate::comments::Comments;
-use ruff_formatter::{write, Buffer, FormatResult};
-use rustpython_parser::ast::ExprCompare;
+use crate::prelude::*;
+use crate::FormatNodeRule;
+use ruff_formatter::{
+    write, FormatError, FormatOwnedWithRule, FormatRefWithRule, FormatRuleWithOptions,
+};
+use ruff_python_ast::prelude::Expr;
+use rustpython_parser::ast::{CmpOp, ExprCompare};
 
 #[derive(Default)]
-pub struct FormatExprCompare;
+pub struct FormatExprCompare {
+    parentheses: Option<Parentheses>,
+}
+
+impl FormatRuleWithOptions<ExprCompare, PyFormatContext<'_>> for FormatExprCompare {
+    type Options = Option<Parentheses>;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.parentheses = options;
+        self
+    }
+}
 
 impl FormatNodeRule<ExprCompare> for FormatExprCompare {
-    fn fmt_fields(&self, _item: &ExprCompare, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(
-            f,
-            [not_yet_implemented_custom_text(
-                "NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right"
-            )]
-        )
+    fn fmt_fields(&self, item: &ExprCompare, f: &mut PyFormatter) -> FormatResult<()> {
+        item.fmt_binary(self.parentheses, f)
+    }
+}
+
+impl<'ast> FormatBinaryLike<'ast> for ExprCompare {
+    type FormatOperator = FormatOwnedWithRule<CmpOp, FormatCmpOp, PyFormatContext<'ast>>;
+
+    fn binary_layout(&self) -> BinaryLayout {
+        if self.ops.len() == 1 {
+            match self.comparators.as_slice() {
+                [right] => BinaryLayout::from_left_right(&self.left, right),
+                [..] => BinaryLayout::Default,
+            }
+        } else {
+            BinaryLayout::Default
+        }
+    }
+
+    fn fmt_default(&self, f: &mut PyFormatter<'ast, '_>) -> FormatResult<()> {
+        let ExprCompare {
+            range: _,
+            left,
+            ops,
+            comparators,
+        } = self;
+
+        let comments = f.context().comments().clone();
+
+        write!(f, [group(&left.format())])?;
+
+        assert_eq!(comparators.len(), ops.len());
+
+        for (operator, comparator) in ops.iter().zip(comparators) {
+            let leading_comparator_comments = comments.leading_comments(comparator);
+            if leading_comparator_comments.is_empty() {
+                write!(f, [soft_line_break_or_space()])?;
+            } else {
+                // Format the expressions leading comments **before** the operator
+                write!(
+                    f,
+                    [
+                        hard_line_break(),
+                        leading_comments(leading_comparator_comments)
+                    ]
+                )?;
+            }
+
+            write!(f, [operator.format(), space(), group(&comparator.format())])?;
+        }
+
+        Ok(())
+    }
+
+    fn left(&self) -> FormatResult<&Expr> {
+        Ok(self.left.as_ref())
+    }
+
+    fn right(&self) -> FormatResult<&Expr> {
+        self.comparators.last().ok_or(FormatError::SyntaxError)
+    }
+
+    fn operator(&self) -> Self::FormatOperator {
+        let op = *self.ops.first().unwrap();
+        op.into_format()
     }
 }
 
@@ -28,6 +101,61 @@ impl NeedsParentheses for ExprCompare {
         source: &str,
         comments: &Comments,
     ) -> Parentheses {
-        default_expression_needs_parentheses(self.into(), parenthesize, source, comments)
+        match default_expression_needs_parentheses(self.into(), parenthesize, source, comments) {
+            parentheses @ Parentheses::Optional => match self.binary_layout() {
+                BinaryLayout::Default => parentheses,
+
+                BinaryLayout::ExpandRight
+                | BinaryLayout::ExpandLeft
+                | BinaryLayout::ExpandRightThenLeft
+                    if self
+                        .comparators
+                        .last()
+                        .map_or(false, |right| comments.has_leading_comments(right)) =>
+                {
+                    parentheses
+                }
+                _ => Parentheses::Custom,
+            },
+            parentheses => parentheses,
+        }
+    }
+}
+
+#[derive(Copy, Clone)]
+pub struct FormatCmpOp;
+
+impl<'ast> AsFormat<PyFormatContext<'ast>> for CmpOp {
+    type Format<'a> = FormatRefWithRule<'a, CmpOp, FormatCmpOp, PyFormatContext<'ast>>;
+
+    fn format(&self) -> Self::Format<'_> {
+        FormatRefWithRule::new(self, FormatCmpOp)
+    }
+}
+
+impl<'ast> IntoFormat<PyFormatContext<'ast>> for CmpOp {
+    type Format = FormatOwnedWithRule<CmpOp, FormatCmpOp, PyFormatContext<'ast>>;
+
+    fn into_format(self) -> Self::Format {
+        FormatOwnedWithRule::new(self, FormatCmpOp)
+    }
+}
+
+impl FormatRule<CmpOp, PyFormatContext<'_>> for FormatCmpOp {
+    fn fmt(&self, item: &CmpOp, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let operator = match item {
+            CmpOp::Eq => "==",
+            CmpOp::NotEq => "!=",
+            CmpOp::Lt => "<",
+            CmpOp::LtE => "<=",
+            CmpOp::Gt => ">",
+            CmpOp::GtE => ">=",
+            CmpOp::Is => "is",
+            CmpOp::IsNot => "is not",
+            CmpOp::In => "in",
+            CmpOp::NotIn => "not in",
+        };
+
+        text(operator).fmt(f)
     }
 }

crates/ruff_python_formatter/src/expression/mod.rs

@@ -76,7 +76,7 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
             Expr::Await(expr) => expr.format().fmt(f),
             Expr::Yield(expr) => expr.format().fmt(f),
             Expr::YieldFrom(expr) => expr.format().fmt(f),
-            Expr::Compare(expr) => expr.format().fmt(f),
+            Expr::Compare(expr) => expr.format().with_options(Some(parentheses)).fmt(f),
             Expr::Call(expr) => expr.format().fmt(f),
             Expr::FormattedValue(expr) => expr.format().fmt(f),
             Expr::JoinedStr(expr) => expr.format().fmt(f),

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__bracketmatch_py.snap

@@ -19,11 +19,9 @@ lambda x=lambda y={1: 3}: y['x':lambda y: {1: 2}]: x
 ```diff
 --- Black
 +++ Ruff
-@@ -1,4 +1,6 @@
+@@ -1,4 +1,4 @@
 -for ((x in {}) or {})["a"] in x:
-+for ((NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right) or {})[
-+    "NOT_YET_IMPLEMENTED_STRING"
-+] in x:
++for ((x in {}) or {})["NOT_YET_IMPLEMENTED_STRING"] in x:
      pass
 -pem_spam = lambda l, spam={"x": 3}: not spam.get(l.strip())
 -lambda x=lambda y={1: 3}: y["x" : lambda y: {1: 2}]: x
@@ -34,9 +32,7 @@ lambda x=lambda y={1: 3}: y['x':lambda y: {1: 2}]: x
 ## Ruff Output
 
 ```py
-for ((NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right) or {})[
-    "NOT_YET_IMPLEMENTED_STRING"
-] in x:
+for ((x in {}) or {})["NOT_YET_IMPLEMENTED_STRING"] in x:
     pass
 pem_spam = lambda x: True
 lambda x: True

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments2_py.snap

@@ -249,7 +249,7 @@ instruction()#comment with bad spacing
  
 -if "PYTHON" in os.environ:
 -    add_compiler(compiler_from_env())
-+if NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right:
++if "NOT_YET_IMPLEMENTED_STRING" in os.environ:
 +    NOT_IMPLEMENTED_call(NOT_IMPLEMENTED_arg)
  else:
      # for compiler in compilers.values():
@@ -283,15 +283,13 @@ instruction()#comment with bad spacing
 +            parameters.children[-1],
 +        ]  # type: ignore
      if (
--        self._proc is not None
-+        NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right
+         self._proc is not None
          # has the child process finished?
--        and self._returncode is None
-+        and NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right
+         and self._returncode is None
          # the child process has finished, but the
          # transport hasn't been notified yet?
 -        and self._proc.poll() is None
-+        and NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right
++        and NOT_IMPLEMENTED_call() is None
      ):
          pass
      # no newline before or after
@@ -349,7 +347,7 @@ instruction()#comment with bad spacing
      while True:
          if False:
              continue
-@@ -141,24 +111,18 @@
+@@ -141,24 +111,19 @@
          # and round and round we go
  
      # let's return
@@ -374,13 +372,14 @@ instruction()#comment with bad spacing
      def _init_host(self, parsed) -> None:
 -        if parsed.hostname is None or not parsed.hostname.strip():  # type: ignore
 +        if (
-+            NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right  # type: ignore
++            parsed.hostname
++            is None  # type: ignore
 +            or not NOT_IMPLEMENTED_call()
 +        ):
              pass
  
  
-@@ -167,7 +131,7 @@
+@@ -167,7 +132,7 @@
  #######################
  
  
@@ -440,7 +439,7 @@ not_shareables = [
     NOT_IMPLEMENTED_call(NOT_IMPLEMENTED_arg),
 ]
 
-if NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right:
+if "NOT_YET_IMPLEMENTED_STRING" in os.environ:
     NOT_IMPLEMENTED_call(NOT_IMPLEMENTED_arg)
 else:
     # for compiler in compilers.values():
@@ -474,12 +473,12 @@ def inline_comments_in_brackets_ruin_everything():
             parameters.children[-1],
         ]  # type: ignore
     if (
-        NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right
+        self._proc is not None
         # has the child process finished?
-        and NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right
+        and self._returncode is None
         # the child process has finished, but the
         # transport hasn't been notified yet?
-        and NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right
+        and NOT_IMPLEMENTED_call() is None
     ):
         pass
     # no newline before or after
@@ -516,7 +515,8 @@ CONFIG_FILES = [CONFIG_FILE] + SHARED_CONFIG_FILES + USER_CONFIG_FILES  # type:
 class Test:
     def _init_host(self, parsed) -> None:
         if (
-            NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right  # type: ignore
+            parsed.hostname
+            is None  # type: ignore
             or not NOT_IMPLEMENTED_call()
         ):
             pass

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments3_py.snap

@@ -83,7 +83,7 @@ def func():
 +    if NOT_IMPLEMENTED_call(NOT_IMPLEMENTED_arg):
          embedded = []
          for exc in exc_value.exceptions:
--            if exc not in _seen:
+             if exc not in _seen:
 -                embedded.append(
 -                    # This should be left alone (before)
 -                    traceback.TracebackException.from_exception(
@@ -97,7 +97,6 @@ def func():
 -                    )
 -                    # This should be left alone (after)
 -                )
-+            if NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right:
 +                NOT_IMPLEMENTED_call(NOT_IMPLEMENTED_arg)
  
      # everything is fine if the expression isn't nested
@@ -130,7 +129,7 @@ def func():
     if NOT_IMPLEMENTED_call(NOT_IMPLEMENTED_arg):
         embedded = []
         for exc in exc_value.exceptions:
-            if NOT_IMPLEMENTED_left < NOT_IMPLEMENTED_right:
+            if exc not in _seen:
                 NOT_IMPLEMENTED_call(NOT_IMPLEMENTED_arg)
 
     # everything is fine if the expression isn't nested