Refactor comparison expression parsing

crates/ruff_python_parser/src/parser/expression.rs

@@ -11,7 +11,6 @@ use ruff_python_ast::{
 };
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
-use crate::parser::helpers::token_kind_to_cmp_op;
 use crate::parser::progress::ParserProgress;
 use crate::parser::{helpers, FunctionKind, Parser};
 use crate::string::{parse_fstring_literal_element, parse_string_literal, StringType};
@@ -357,21 +356,16 @@ impl<'src> Parser<'src> {
                 continue;
             }
 
-            // Operator token.
-            self.bump(token);
-
-            if token.is_compare_operator() {
-                left = Expr::Compare(self.parse_compare_expression(
-                    left.expr,
-                    start,
-                    token,
-                    new_precedence,
-                    context,
-                ))
-                .into();
+            if let Some(cmp_op) = token.as_compare_operator(self.peek()) {
+                left =
+                    Expr::Compare(self.parse_compare_expression(left.expr, start, cmp_op, context))
+                        .into();
                 continue;
             }
 
+            // Operator token.
+            self.bump(token);
+
             let right = self.parse_binary_expression_or_higher(new_precedence, context);
 
             left.expr = Expr::BinOp(ast::ExprBinOp {
@@ -1079,6 +1073,27 @@ impl<'src> Parser<'src> {
         }
     }
 
+    /// Bump the appropriate token(s) for the given comparison operator.
+    fn bump_cmp_op(&mut self, op: CmpOp) {
+        let (first, second) = match op {
+            CmpOp::Eq => (TokenKind::EqEqual, None),
+            CmpOp::NotEq => (TokenKind::NotEqual, None),
+            CmpOp::Lt => (TokenKind::Less, None),
+            CmpOp::LtE => (TokenKind::LessEqual, None),
+            CmpOp::Gt => (TokenKind::Greater, None),
+            CmpOp::GtE => (TokenKind::GreaterEqual, None),
+            CmpOp::Is => (TokenKind::Is, None),
+            CmpOp::IsNot => (TokenKind::Is, Some(TokenKind::Not)),
+            CmpOp::In => (TokenKind::In, None),
+            CmpOp::NotIn => (TokenKind::Not, Some(TokenKind::In)),
+        };
+
+        self.bump(first);
+        if let Some(second) = second {
+            self.bump(second);
+        }
+    }
+
     /// Parse a comparison expression.
     ///
     /// This includes the following operators:
@@ -1095,68 +1110,43 @@ impl<'src> Parser<'src> {
         &mut self,
         lhs: Expr,
         start: TextSize,
-        operator: TokenKind,
-        operator_binding_power: OperatorPrecedence,
+        op: CmpOp,
         context: ExpressionContext,
     ) -> ast::ExprCompare {
-        let compare_operator = token_kind_to_cmp_op([operator, self.current_token_kind()]).unwrap();
-
-        // Bump the appropriate token when the compare operator is made up of
-        // two separate tokens.
-        match compare_operator {
-            CmpOp::IsNot => {
-                self.bump(TokenKind::Not);
-            }
-            CmpOp::NotIn => {
-                self.bump(TokenKind::In);
-            }
-            _ => {}
-        }
+        self.bump_cmp_op(op);
 
         let mut comparators = vec![];
-        let mut compare_operators = vec![compare_operator];
+        let mut operators = vec![op];
 
         let mut progress = ParserProgress::default();
 
         loop {
             progress.assert_progressing(self);
 
-            let parsed_expr =
-                self.parse_binary_expression_or_higher(operator_binding_power, context);
-            comparators.push(parsed_expr.expr);
+            comparators.push(
+                self.parse_binary_expression_or_higher(
+                    OperatorPrecedence::ComparisonsMembershipIdentity,
+                    context,
+                )
+                .expr,
+            );
 
-            let next_operator = self.current_token_kind();
-            if !next_operator.is_compare_operator()
-                || (matches!(next_operator, TokenKind::In) && !context.is_in_included())
-            {
+            let next_token = self.current_token_kind();
+            if matches!(next_token, TokenKind::In) && !context.is_in_included() {
                 break;
             }
-            self.bump(next_operator); // compare operator
 
-            if let Ok(compare_operator) =
-                token_kind_to_cmp_op([next_operator, self.current_token_kind()])
-            {
-                // Bump the appropriate token when the compare operator is made up of
-                // two separate tokens.
-                match compare_operator {
-                    CmpOp::IsNot => {
-                        self.bump(TokenKind::Not);
-                    }
-                    CmpOp::NotIn => {
-                        self.bump(TokenKind::In);
-                    }
-                    _ => {}
-                }
-
-                compare_operators.push(compare_operator);
-            } else {
+            let Some(next_op) = next_token.as_compare_operator(self.peek()) else {
                 break;
-            }
+            };
+
+            self.bump_cmp_op(next_op);
+            operators.push(next_op);
         }
 
         ast::ExprCompare {
             left: Box::new(lhs),
-            ops: compare_operators.into_boxed_slice(),
+            ops: operators.into_boxed_slice(),
             comparators: comparators.into_boxed_slice(),
             range: self.node_range(start),
         }

crates/ruff_python_parser/src/parser/helpers.rs

@@ -1,6 +1,4 @@
-use ruff_python_ast::{self as ast, CmpOp, Expr, ExprContext};
-
-use crate::TokenKind;
+use ruff_python_ast::{self as ast, Expr, ExprContext};
 
 /// Set the `ctx` for `Expr::Id`, `Expr::Attribute`, `Expr::Subscript`, `Expr::Starred`,
 /// `Expr::Tuple` and `Expr::List`. If `expr` is either `Expr::Tuple` or `Expr::List`,
@@ -26,20 +24,3 @@ pub(super) fn set_expr_ctx(expr: &mut Expr, new_ctx: ExprContext) {
         _ => {}
     }
 }
-
-/// Converts a [`TokenKind`] array of size 2 to its correspondent [`CmpOp`].
-pub(super) fn token_kind_to_cmp_op(kind: [TokenKind; 2]) -> Result<CmpOp, ()> {
-    Ok(match kind {
-        [TokenKind::Is, TokenKind::Not] => CmpOp::IsNot,
-        [TokenKind::Is, _] => CmpOp::Is,
-        [TokenKind::Not, TokenKind::In] => CmpOp::NotIn,
-        [TokenKind::In, _] => CmpOp::In,
-        [TokenKind::EqEqual, _] => CmpOp::Eq,
-        [TokenKind::NotEqual, _] => CmpOp::NotEq,
-        [TokenKind::Less, _] => CmpOp::Lt,
-        [TokenKind::LessEqual, _] => CmpOp::LtE,
-        [TokenKind::Greater, _] => CmpOp::Gt,
-        [TokenKind::GreaterEqual, _] => CmpOp::GtE,
-        _ => return Err(()),
-    })
-}

crates/ruff_python_parser/src/token.rs

@@ -5,7 +5,7 @@
 //!
 //! [CPython source]: https://github.com/python/cpython/blob/dfc2e065a2e71011017077e549cd2f9bf4944c54/Include/internal/pycore_token.h;
 
-use ruff_python_ast::{AnyStringKind, BoolOp, Int, IpyEscapeKind, Operator, UnaryOp};
+use ruff_python_ast::{AnyStringKind, BoolOp, CmpOp, Int, IpyEscapeKind, Operator, UnaryOp};
 use std::fmt;
 
 use crate::Mode;
@@ -699,20 +699,31 @@ impl TokenKind {
         matches!(self, TokenKind::Match | TokenKind::Case)
     }
 
+    /// Returns `true` if the current token is a comparison operator. This function requires the
+    /// next token operators containing two tokens such as `not in` and `is not`.
     #[inline]
-    pub const fn is_compare_operator(&self) -> bool {
-        matches!(
-            self,
-            TokenKind::Not
-                | TokenKind::In
-                | TokenKind::Is
-                | TokenKind::EqEqual
-                | TokenKind::NotEqual
-                | TokenKind::Less
-                | TokenKind::LessEqual
-                | TokenKind::Greater
-                | TokenKind::GreaterEqual
-        )
+    pub const fn is_compare_operator(&self, next: TokenKind) -> bool {
+        self.as_compare_operator(next).is_some()
+    }
+
+    /// Returns the [`CmpOp`] that corresponds to this token kind, if it is a comparison operator,
+    /// otherwise return [None]. This function requires the next token operators containing two
+    /// tokens such as `not in` and `is not`.
+    #[inline]
+    pub const fn as_compare_operator(&self, next: TokenKind) -> Option<CmpOp> {
+        Some(match self {
+            TokenKind::In => CmpOp::In,
+            TokenKind::Not if matches!(next, TokenKind::In) => CmpOp::NotIn,
+            TokenKind::Is if matches!(next, TokenKind::Not) => CmpOp::IsNot,
+            TokenKind::Is => CmpOp::Is,
+            TokenKind::EqEqual => CmpOp::Eq,
+            TokenKind::NotEqual => CmpOp::NotEq,
+            TokenKind::Less => CmpOp::Lt,
+            TokenKind::LessEqual => CmpOp::LtE,
+            TokenKind::Greater => CmpOp::Gt,
+            TokenKind::GreaterEqual => CmpOp::GtE,
+            _ => return None,
+        })
     }
 
     /// Returns `true` if the current token is a boolean operator.