pyflakes F632 Autofix (#612)

README.md

@@ -322,7 +322,7 @@ For more, see [Pyflakes](https://pypi.org/project/pyflakes/2.5.0/) on PyPI.
 | F621 | ExpressionsInStarAssignment | Too many expressions in star-unpacking assignment |  |
 | F622 | TwoStarredExpressions | Two starred expressions in assignment |  |
 | F631 | AssertTuple | Assert test is a non-empty tuple, which is always `True` |  |
-| F632 | IsLiteral | Use `==` and `!=` to compare constant literals |  |
+| F632 | IsLiteral | Use `==` and `!=` to compare constant literals | 🛠 |
 | F633 | InvalidPrintSyntax | Use of `>>` is invalid with `print` function |  |
 | F634 | IfTuple | If test is a tuple, which is always `True` |  |
 | F701 | BreakOutsideLoop | `break` outside loop |  |

resources/test/fixtures/F632.py

@@ -3,3 +3,6 @@ if x is "abc":
 
 if 123 is not y:
     pass
+
+if "123" is x < 3:
+    pass

src/check_ast.rs

@@ -1324,12 +1324,13 @@ where
                 }
 
                 if self.settings.enabled.contains(&CheckCode::F632) {
-                    self.checks.extend(pyflakes::checks::is_literal(
+                    pyflakes::plugins::invalid_literal_comparison(
+                        self,
                         left,
                         ops,
                         comparators,
                         self.locate_check(Range::from_located(expr)),
-                    ));
+                    );
                 }
 
                 if self.settings.enabled.contains(&CheckCode::E721) {

src/checks.rs

@@ -1583,6 +1583,7 @@ impl CheckKind {
                 | CheckKind::UsePEP604Annotation
                 | CheckKind::UselessMetaclassType
                 | CheckKind::UselessObjectInheritance(_)
+                | CheckKind::IsLiteral
         )
     }
 }

src/cst/matchers.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use libcst_native::Module;
+use libcst_native::{Expr, Module, SmallStatement, Statement};
 
 pub fn match_module(module_text: &str) -> Result<Module> {
     match libcst_native::parse_module(module_text, None) {
@@ -7,3 +7,15 @@ pub fn match_module(module_text: &str) -> Result<Module> {
         Err(_) => Err(anyhow::anyhow!("Failed to extract CST from source.")),
     }
 }
+
+pub fn match_expr<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut Expr<'b>> {
+    if let Some(Statement::Simple(expr)) = module.body.first_mut() {
+        if let Some(SmallStatement::Expr(expr)) = expr.body.first_mut() {
+            Ok(expr)
+        } else {
+            Err(anyhow::anyhow!("Expected node to be: SmallStatement::Expr"))
+        }
+    } else {
+        Err(anyhow::anyhow!("Expected node to be: Statement::Simple"))
+    }
+}

src/flake8_comprehensions/fixes.rs

@@ -1,28 +1,15 @@
 use anyhow::Result;
 use libcst_native::{
     Arg, Call, Codegen, Dict, DictComp, DictElement, Element, Expr, Expression, LeftCurlyBrace,
-    LeftParen, LeftSquareBracket, List, ListComp, Module, Name, ParenthesizableWhitespace,
-    RightCurlyBrace, RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace,
-    SmallStatement, Statement, Tuple,
+    LeftParen, LeftSquareBracket, List, ListComp, Name, ParenthesizableWhitespace, RightCurlyBrace,
+    RightParen, RightSquareBracket, Set, SetComp, SimpleString, SimpleWhitespace, Tuple,
 };
 
 use crate::ast::types::Range;
 use crate::autofix::Fix;
-use crate::cst::matchers::match_module;
+use crate::cst::matchers::{match_expr, match_module};
 use crate::source_code_locator::SourceCodeLocator;
 
-fn match_expr<'a, 'b>(module: &'a mut Module<'b>) -> Result<&'a mut Expr<'b>> {
-    if let Some(Statement::Simple(expr)) = module.body.first_mut() {
-        if let Some(SmallStatement::Expr(expr)) = expr.body.first_mut() {
-            Ok(expr)
-        } else {
-            Err(anyhow::anyhow!("Expected node to be: SmallStatement::Expr"))
-        }
-    } else {
-        Err(anyhow::anyhow!("Expected node to be: Statement::Simple"))
-    }
-}
-
 fn match_call<'a, 'b>(expr: &'a mut Expr<'b>) -> Result<&'a mut Call<'b>> {
     if let Expression::Call(call) = &mut expr.value {
         Ok(call)

src/pyflakes/checks.rs

@@ -1,10 +1,8 @@
 use std::collections::BTreeSet;
 
-use itertools::izip;
 use regex::Regex;
 use rustpython_parser::ast::{
-    Arg, Arguments, Cmpop, Constant, Excepthandler, ExcepthandlerKind, Expr, ExprKind, Stmt,
-    StmtKind,
+    Arg, Arguments, Constant, Excepthandler, ExcepthandlerKind, Expr, ExprKind, Stmt, StmtKind,
 };
 
 use crate::ast::types::{BindingKind, CheckLocator, FunctionScope, Range, Scope, ScopeKind};
@@ -166,45 +164,6 @@ pub fn repeated_keys(
     checks
 }
 
-fn is_constant(expr: &Expr) -> bool {
-    match &expr.node {
-        ExprKind::Constant { .. } => true,
-        ExprKind::Tuple { elts, .. } => elts.iter().all(is_constant),
-        _ => false,
-    }
-}
-
-fn is_singleton(expr: &Expr) -> bool {
-    matches!(
-        expr.node,
-        ExprKind::Constant {
-            value: Constant::None | Constant::Bool(_) | Constant::Ellipsis,
-            ..
-        }
-    )
-}
-
-fn is_constant_non_singleton(expr: &Expr) -> bool {
-    is_constant(expr) && !is_singleton(expr)
-}
-
-/// F632
-pub fn is_literal(left: &Expr, ops: &[Cmpop], comparators: &[Expr], location: Range) -> Vec<Check> {
-    let mut checks: Vec<Check> = vec![];
-
-    let mut left = left;
-    for (op, right) in izip!(ops, comparators) {
-        if matches!(op, Cmpop::Is | Cmpop::IsNot)
-            && (is_constant_non_singleton(left) || is_constant_non_singleton(right))
-        {
-            checks.push(Check::new(CheckKind::IsLiteral, location));
-        }
-        left = right;
-    }
-
-    checks
-}
-
 /// F621, F622
 pub fn starred_expressions(
     elts: &[Expr],

src/pyflakes/fixes.rs

@@ -1,11 +1,14 @@
 use anyhow::Result;
-use libcst_native::{Codegen, ImportNames, NameOrAttribute, SmallStatement, Statement};
+use libcst_native::{
+    Codegen, CompOp, Comparison, ComparisonTarget, Expr, Expression, ImportNames, NameOrAttribute,
+    SmallStatement, Statement,
+};
 use rustpython_ast::Stmt;
 
 use crate::ast::types::Range;
 use crate::autofix::{helpers, Fix};
 use crate::cst::helpers::compose_module_path;
-use crate::cst::matchers::match_module;
+use crate::cst::matchers::{match_expr, match_module};
 use crate::source_code_locator::SourceCodeLocator;
 
 /// Generate a Fix to remove any unused imports from an `import` statement.
@@ -138,3 +141,66 @@ pub fn remove_unused_import_froms(
         ))
     }
 }
+
+fn match_comparison<'a, 'b>(expr: &'a mut Expr<'b>) -> Result<&'a mut Comparison<'b>> {
+    if let Expression::Comparison(comparison) = &mut expr.value {
+        Ok(comparison)
+    } else {
+        Err(anyhow::anyhow!(
+            "Expected node to be: Expression::Comparison"
+        ))
+    }
+}
+
+/// Generate a Fix to replace invalid is/is not comparisons with equal/not equal
+pub fn fix_invalid_literal_comparison(locator: &SourceCodeLocator, location: Range) -> Result<Fix> {
+    let module_text = locator.slice_source_code_range(&location);
+    let mut tree = match_module(&module_text)?;
+    let mut expr = match_expr(&mut tree)?;
+    let cmp = match_comparison(expr)?;
+    let target = cmp
+        .comparisons
+        .get(0)
+        .ok_or_else(|| anyhow::anyhow!("Expected one ComparisonTarget"))?;
+
+    let new_operator = match &target.operator {
+        CompOp::Is {
+            whitespace_before: b,
+            whitespace_after: a,
+        } => Ok(CompOp::Equal {
+            whitespace_before: b.clone(),
+            whitespace_after: a.clone(),
+        }),
+        CompOp::IsNot {
+            whitespace_before: b,
+            whitespace_after: a,
+            whitespace_between: _,
+        } => Ok(CompOp::NotEqual {
+            whitespace_before: b.clone(),
+            whitespace_after: a.clone(),
+        }),
+        op => Err(anyhow::anyhow!(
+            "Unexpected operator: {:?}. Expected Is or IsNot.",
+            op
+        )),
+    }?;
+
+    expr.value = Expression::Comparison(Box::new(Comparison {
+        left: cmp.left.clone(),
+        comparisons: vec![ComparisonTarget {
+            operator: new_operator,
+            comparator: target.comparator.clone(),
+        }],
+        lpar: cmp.lpar.clone(),
+        rpar: cmp.rpar.clone(),
+    }));
+
+    let mut state = Default::default();
+    tree.codegen(&mut state);
+
+    Ok(Fix::replacement(
+        state.to_string(),
+        location.location,
+        location.end_location,
+    ))
+}

src/pyflakes/plugins/invalid_literal_comparisons.rs

@@ -0,0 +1,62 @@
+use itertools::izip;
+use log::error;
+use rustpython_ast::{Cmpop, Constant, Expr, ExprKind};
+
+use crate::ast::types::Range;
+use crate::check_ast::Checker;
+use crate::checks::{Check, CheckKind};
+use crate::pyflakes::fixes::fix_invalid_literal_comparison;
+
+fn is_singleton(expr: &Expr) -> bool {
+    matches!(
+        expr.node,
+        ExprKind::Constant {
+            value: Constant::None | Constant::Bool(_) | Constant::Ellipsis,
+            ..
+        }
+    )
+}
+
+fn is_constant(expr: &Expr) -> bool {
+    match &expr.node {
+        ExprKind::Constant { .. } => true,
+        ExprKind::Tuple { elts, .. } => elts.iter().all(is_constant),
+        _ => false,
+    }
+}
+
+fn is_constant_non_singleton(expr: &Expr) -> bool {
+    is_constant(expr) && !is_singleton(expr)
+}
+
+/// F632
+pub fn invalid_literal_comparison(
+    checker: &mut Checker,
+    left: &Expr,
+    ops: &[Cmpop],
+    comparators: &[Expr],
+    location: Range,
+) {
+    let mut left = left;
+    for (op, right) in izip!(ops, comparators) {
+        if matches!(op, Cmpop::Is | Cmpop::IsNot)
+            && (is_constant_non_singleton(left) || is_constant_non_singleton(right))
+        {
+            let mut check = Check::new(CheckKind::IsLiteral, location);
+            if checker.patch() {
+                match fix_invalid_literal_comparison(
+                    checker.locator,
+                    Range {
+                        location: left.location,
+                        end_location: right.end_location.unwrap(),
+                    },
+                ) {
+                    Ok(fix) => check.amend(fix),
+                    Err(e) => error!("Failed to fix invalid comparison: {}", e),
+                }
+            }
+            checker.add_check(check);
+        }
+        left = right;
+    }
+}

src/pyflakes/plugins/mod.rs

@@ -1,9 +1,11 @@
 pub use assert_tuple::assert_tuple;
 pub use if_tuple::if_tuple;
+pub use invalid_literal_comparisons::invalid_literal_comparison;
 pub use invalid_print_syntax::invalid_print_syntax;
 pub use raise_not_implemented::raise_not_implemented;
 
 mod assert_tuple;
 mod if_tuple;
+mod invalid_literal_comparisons;
 mod invalid_print_syntax;
 mod raise_not_implemented;

src/snapshots/ruff__linter__tests__F632_F632.py.snap

@@ -9,7 +9,16 @@ expression: checks
   end_location:
     row: 1
     column: 13
-  fix: ~
+  fix:
+    patch:
+      content: "x == \"abc\""
+      location:
+        row: 1
+        column: 3
+      end_location:
+        row: 1
+        column: 13
+    applied: false
 - kind: IsLiteral
   location:
     row: 4
@@ -17,5 +26,31 @@ expression: checks
   end_location:
     row: 4
     column: 15
-  fix: ~
+  fix:
+    patch:
+      content: 123 != y
+      location:
+        row: 4
+        column: 3
+      end_location:
+        row: 4
+        column: 15
+    applied: false
+- kind: IsLiteral
+  location:
+    row: 7
+    column: 9
+  end_location:
+    row: 7
+    column: 17
+  fix:
+    patch:
+      content: "\"123\" == x"
+      location:
+        row: 7
+        column: 3
+      end_location:
+        row: 7
+        column: 13
+    applied: false