Avoid removing statements that contain side-effects (#1920)

Closes #1917.
astral-sh · Jan 16, 2023 · f3bf008 · f3bf008
1 parent 3b4aaa5
commit f3bf008
Show file tree

Hide file tree

Showing 6 changed files with 114 additions and 76 deletions.
diff --git a/resources/test/fixtures/pyflakes/F841_3.py b/resources/test/fixtures/pyflakes/F841_3.py
@@ -46,6 +46,17 @@ def f():
         pass
 
 
+def f(a, b):
+    x = (
+        a()
+        if a is not None
+        else b
+    )
+
+    y = \
+        a() if a is not None else b
+
+
 def f(a, b):
     x = (
         a

diff --git a/src/ast/helpers.rs b/src/ast/helpers.rs
@@ -94,6 +94,45 @@ pub fn contains_call_path(checker: &Checker, expr: &Expr, target: &[&str]) -> bo
     })
 }
 
+/// Return `true` if the `Expr` contains an expression that appears to include a
+/// side-effect (like a function call).
+pub fn contains_effect(checker: &Checker, expr: &Expr) -> bool {
+    any_over_expr(expr, &|expr| {
+        // Accept empty initializers.
+        if let ExprKind::Call {
+            func,
+            args,
+            keywords,
+        } = &expr.node
+        {
+            if args.is_empty() && keywords.is_empty() {
+                if let ExprKind::Name { id, .. } = &func.node {
+                    let is_empty_initializer = (id == "set"
+                        || id == "list"
+                        || id == "tuple"
+                        || id == "dict"
+                        || id == "frozenset")
+                        && checker.is_builtin(id);
+                    return !is_empty_initializer;
+                }
+            }
+        }
+
+        // Otherwise, avoid all complex expressions.
+        matches!(
+            expr.node,
+            ExprKind::Call { .. }
+                | ExprKind::Await { .. }
+                | ExprKind::GeneratorExp { .. }
+                | ExprKind::ListComp { .. }
+                | ExprKind::SetComp { .. }
+                | ExprKind::DictComp { .. }
+                | ExprKind::Yield { .. }
+                | ExprKind::YieldFrom { .. }
+        )
+    })
+}
+
 /// Call `func` over every `Expr` in `expr`, returning `true` if any expression
 /// returns `true`..
 pub fn any_over_expr<F>(expr: &Expr, func: &F) -> bool

diff --git a/src/rules/flake8_simplify/rules/ast_if.rs b/src/rules/flake8_simplify/rules/ast_if.rs
@@ -2,7 +2,7 @@ use rustpython_ast::{Cmpop, Constant, Expr, ExprContext, ExprKind, Stmt, StmtKin
 
 use crate::ast::comparable::ComparableExpr;
 use crate::ast::helpers::{
-    any_over_expr, contains_call_path, create_expr, create_stmt, has_comments, unparse_expr,
+    contains_call_path, contains_effect, create_expr, create_stmt, has_comments, unparse_expr,
     unparse_stmt,
 };
 use crate::ast::types::Range;
@@ -272,19 +272,7 @@ pub fn use_dict_get_with_default(
     }
 
     // Check that the default value is not "complex".
-    if any_over_expr(default_val, &|expr| {
-        matches!(
-            expr.node,
-            ExprKind::Call { .. }
-                | ExprKind::Await { .. }
-                | ExprKind::GeneratorExp { .. }
-                | ExprKind::ListComp { .. }
-                | ExprKind::SetComp { .. }
-                | ExprKind::DictComp { .. }
-                | ExprKind::Yield { .. }
-                | ExprKind::YieldFrom { .. }
-        )
-    }) {
+    if contains_effect(checker, default_val) {
         return;
     }
 

diff --git a/src/rules/pyflakes/rules/unused_variable.rs b/src/rules/pyflakes/rules/unused_variable.rs
@@ -1,9 +1,10 @@
 use itertools::Itertools;
 use log::error;
-use rustpython_ast::{Expr, ExprKind, Location, Stmt, StmtKind};
+use rustpython_ast::{ExprKind, Location, Stmt, StmtKind};
 use rustpython_parser::lexer;
 use rustpython_parser::lexer::Tok;
 
+use crate::ast::helpers::contains_effect;
 use crate::ast::types::{BindingKind, Range, RefEquality, ScopeKind};
 use crate::autofix::helpers::delete_stmt;
 use crate::checkers::ast::Checker;
@@ -12,41 +13,6 @@ use crate::registry::{Diagnostic, RuleCode};
 use crate::source_code::Locator;
 use crate::violations;
 
-fn is_literal_or_name(expr: &Expr, checker: &Checker) -> bool {
-    // Accept any obvious literals or names.
-    if matches!(
-        expr.node,
-        ExprKind::Constant { .. }
-            | ExprKind::Name { .. }
-            | ExprKind::List { .. }
-            | ExprKind::Tuple { .. }
-            | ExprKind::Set { .. }
-    ) {
-        return true;
-    }
-
-    // Accept empty initializers.
-    if let ExprKind::Call {
-        func,
-        args,
-        keywords,
-    } = &expr.node
-    {
-        if args.is_empty() && keywords.is_empty() {
-            if let ExprKind::Name { id, .. } = &func.node {
-                return (id == "set"
-                    || id == "list"
-                    || id == "tuple"
-                    || id == "dict"
-                    || id == "frozenset")
-                    && checker.is_builtin(id);
-            }
-        }
-    }
-
-    false
-}
-
 fn match_token_after<F>(stmt: &Stmt, locator: &Locator, f: F) -> Location
 where
     F: Fn(Tok) -> bool,
@@ -78,8 +44,18 @@ fn remove_unused_variable(
     // First case: simple assignment (`x = 1`)
     if let StmtKind::Assign { targets, value, .. } = &stmt.node {
         if targets.len() == 1 && matches!(targets[0].node, ExprKind::Name { .. }) {
-            return if is_literal_or_name(value, checker) {
-                // If assigning to a constant (`x = 1`), delete the entire statement.
+            return if contains_effect(checker, value) {
+                // If the expression is complex (`x = foo()`), remove the assignment,
+                // but preserve the right-hand side.
+                Some((
+                    DeletionKind::Partial,
+                    Fix::deletion(
+                        stmt.location,
+                        match_token_after(stmt, checker.locator, |tok| tok == Tok::Equal),
+                    ),
+                ))
+            } else {
+                // If (e.g.) assigning to a constant (`x = 1`), delete the entire statement.
                 let parent = checker
                     .child_to_parent
                     .get(&RefEquality(stmt))
@@ -98,16 +74,6 @@ fn remove_unused_variable(
                         None
                     }
                 }
-            } else {
-                // If the expression is more complex (`x = foo()`), remove the assignment,
-                // but preserve the right-hand side.
-                Some((
-                    DeletionKind::Partial,
-                    Fix::deletion(
-                        stmt.location,
-                        match_token_after(stmt, checker.locator, |tok| tok == Tok::Equal),
-                    ),
-                ))
             };
         }
     }
@@ -120,7 +86,17 @@ fn remove_unused_variable(
     } = &stmt.node
     {
         if matches!(target.node, ExprKind::Name { .. }) {
-            return if is_literal_or_name(value, checker) {
+            return if contains_effect(checker, value) {
+                // If the expression is complex (`x = foo()`), remove the assignment,
+                // but preserve the right-hand side.
+                Some((
+                    DeletionKind::Partial,
+                    Fix::deletion(
+                        stmt.location,
+                        match_token_after(stmt, checker.locator, |tok| tok == Tok::Equal),
+                    ),
+                ))
+            } else {
                 // If assigning to a constant (`x = 1`), delete the entire statement.
                 let parent = checker
                     .child_to_parent
@@ -140,16 +116,6 @@ fn remove_unused_variable(
                         None
                     }
                 }
-            } else {
-                // If the expression is more complex (`x = foo()`), remove the assignment,
-                // but preserve the right-hand side.
-                Some((
-                    DeletionKind::Partial,
-                    Fix::deletion(
-                        stmt.location,
-                        match_token_after(stmt, checker.locator, |tok| tok == Tok::Equal),
-                    ),
-                ))
             };
         }
     }

diff --git a/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F841_F841_0.py.snap b/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F841_F841_0.py.snap
@@ -31,10 +31,10 @@ expression: diagnostics
     content: ""
     location:
       row: 16
-      column: 4
+      column: 0
     end_location:
-      row: 16
-      column: 8
+      row: 17
+      column: 0
   parent: ~
 - kind:
     UnusedVariable: foo

diff --git a/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F841_F841_3.py.snap b/src/rules/pyflakes/snapshots/ruff__rules__pyflakes__tests__F841_F841_3.py.snap
@@ -246,4 +246,38 @@ expression: diagnostics
       row: 57
       column: 8
   parent: ~
+- kind:
+    UnusedVariable: x
+  location:
+    row: 61
+    column: 4
+  end_location:
+    row: 61
+    column: 5
+  fix:
+    content: pass
+    location:
+      row: 61
+      column: 4
+    end_location:
+      row: 65
+      column: 5
+  parent: ~
+- kind:
+    UnusedVariable: y
+  location:
+    row: 67
+    column: 4
+  end_location:
+    row: 67
+    column: 5
+  fix:
+    content: ""
+    location:
+      row: 67
+      column: 0
+    end_location:
+      row: 69
+      column: 0
+  parent: ~