Parenthesize expressions prior to lexing in F632 (#5001)

crates/ruff_python_ast/src/helpers.rs

@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::ops::Sub;
 use std::path::Path;
 
 use itertools::Itertools;
@@ -1444,15 +1445,18 @@ impl LocatedCmpop {
 /// `RustPython` doesn't include line and column information on [`Cmpop`] nodes.
 /// `CPython` doesn't either. This method iterates over the token stream and
 /// re-identifies [`Cmpop`] nodes, annotating them with valid ranges.
-pub fn locate_cmpops(contents: &str) -> Vec<LocatedCmpop> {
-    let mut tok_iter = lexer::lex(contents, Mode::Expression)
+pub fn locate_cmpops(expr: &Expr, locator: &Locator) -> Vec<LocatedCmpop> {
+    // If `Expr` is a multi-line expression, we need to parenthesize it to
+    // ensure that it's lexed correctly.
+    let contents = locator.slice(expr.range());
+    let parenthesized_contents = format!("({contents})");
+    let mut tok_iter = lexer::lex(&parenthesized_contents, Mode::Expression)
         .flatten()
-        .filter(|(tok, _)|
-            // Skip whitespace, including Tok::Newline. It should be sufficient to skip
-            // Tok::NonLogicalNewline, but the lexer doesn't know that we're within an
-            // expression, and so it may confusion logical and non-logical newlines.
-            !matches!(tok, Tok::Newline | Tok::NonLogicalNewline | Tok::Comment(_)))
+        .skip(1)
+        .map(|(tok, range)| (tok, range.sub(TextSize::from(1))))
+        .filter(|(tok, _)| !matches!(tok, Tok::NonLogicalNewline | Tok::Comment(_)))
         .peekable();
+
     let mut ops: Vec<LocatedCmpop> = vec![];
     let mut count = 0u32;
     loop {
@@ -1617,7 +1621,7 @@ mod tests {
 
     use anyhow::Result;
     use ruff_text_size::{TextLen, TextRange, TextSize};
-    use rustpython_ast::{Stmt, Suite};
+    use rustpython_ast::{Expr, Stmt, Suite};
     use rustpython_parser::ast::Cmpop;
     use rustpython_parser::Parse;
 
@@ -1808,13 +1812,11 @@ else:
 
     #[test]
     fn extract_elif_else_range() -> Result<()> {
-        let contents = "
-if a:
+        let contents = "if a:
     ...
 elif b:
     ...
-"
-        .trim_start();
+";
         let program = Suite::parse(contents, "<filename>")?;
         let stmt = program.first().unwrap();
         let stmt = Stmt::as_if_stmt(stmt).unwrap();
@@ -1823,13 +1825,11 @@ elif b:
         assert_eq!(range.start(), TextSize::from(14));
         assert_eq!(range.end(), TextSize::from(18));
 
-        let contents = "
-if a:
+        let contents = "if a:
     ...
 else:
     ...
-"
-        .trim_start();
+";
         let program = Suite::parse(contents, "<filename>")?;
         let stmt = program.first().unwrap();
         let stmt = Stmt::as_if_stmt(stmt).unwrap();
@@ -1842,61 +1842,84 @@ else:
     }
 
     #[test]
-    fn extract_cmpop_location() {
+    fn extract_cmpop_location() -> Result<()> {
+        let contents = "x == 1";
+        let expr = Expr::parse(contents, "<filename>")?;
+        let locator = Locator::new(contents);
         assert_eq!(
-            locate_cmpops("x == 1"),
+            locate_cmpops(&expr, &locator),
             vec![LocatedCmpop::new(
                 TextSize::from(2)..TextSize::from(4),
                 Cmpop::Eq
             )]
         );
 
+        let contents = "x != 1";
+        let expr = Expr::parse(contents, "<filename>")?;
+        let locator = Locator::new(contents);
         assert_eq!(
-            locate_cmpops("x != 1"),
+            locate_cmpops(&expr, &locator),
             vec![LocatedCmpop::new(
                 TextSize::from(2)..TextSize::from(4),
                 Cmpop::NotEq
             )]
         );
 
+        let contents = "x is 1";
+        let expr = Expr::parse(contents, "<filename>")?;
+        let locator = Locator::new(contents);
         assert_eq!(
-            locate_cmpops("x is 1"),
+            locate_cmpops(&expr, &locator),
             vec![LocatedCmpop::new(
                 TextSize::from(2)..TextSize::from(4),
                 Cmpop::Is
             )]
         );
 
+        let contents = "x is not 1";
+        let expr = Expr::parse(contents, "<filename>")?;
+        let locator = Locator::new(contents);
         assert_eq!(
-            locate_cmpops("x is not 1"),
+            locate_cmpops(&expr, &locator),
             vec![LocatedCmpop::new(
                 TextSize::from(2)..TextSize::from(8),
                 Cmpop::IsNot
             )]
         );
 
+        let contents = "x in 1";
+        let expr = Expr::parse(contents, "<filename>")?;
+        let locator = Locator::new(contents);
         assert_eq!(
-            locate_cmpops("x in 1"),
+            locate_cmpops(&expr, &locator),
             vec![LocatedCmpop::new(
                 TextSize::from(2)..TextSize::from(4),
                 Cmpop::In
             )]
         );
 
+        let contents = "x not in 1";
+        let expr = Expr::parse(contents, "<filename>")?;
+        let locator = Locator::new(contents);
         assert_eq!(
-            locate_cmpops("x not in 1"),
+            locate_cmpops(&expr, &locator),
             vec![LocatedCmpop::new(
                 TextSize::from(2)..TextSize::from(8),
                 Cmpop::NotIn
             )]
         );
 
+        let contents = "x != (1 is not 2)";
+        let expr = Expr::parse(contents, "<filename>")?;
+        let locator = Locator::new(contents);
         assert_eq!(
-            locate_cmpops("x != (1 is not 2)"),
+            locate_cmpops(&expr, &locator),
             vec![LocatedCmpop::new(
                 TextSize::from(2)..TextSize::from(4),
                 Cmpop::NotEq
             )]
         );
+
+        Ok(())
     }
 }