impl remove_unused_variable #2207

pyrefly/lib/state/lsp.rs

@@ -3938,6 +3938,22 @@ impl<'a> Transaction<'a> {
                 _ => {}
             }
         }
+        if let Some(answers) = self.get_answers(handle) {
+            let bindings = answers.bindings();
+            for unused in bindings.unused_variables() {
+                if unused.range.contains_range(range)
+                    && let Some(action) =
+                        quick_fixes::remove_unused_variable::remove_unused_variable_code_action(
+                            &module_info,
+                            &ast,
+                            unused.name.as_str(),
+                            unused.range,
+                        )
+                {
+                    other_actions.push(action);
+                }
+            }
+        }
 
         import_actions.sort();
 

pyrefly/lib/state/lsp/quick_fixes.rs

@@ -29,6 +29,7 @@ pub(crate) mod move_module;
 pub(crate) mod pyrefly_ignore;
 pub(crate) mod pytest_fixture;
 pub(crate) mod redundant_cast;
+pub(crate) mod remove_unused_variable;
 pub(crate) mod safe_delete;
 pub(crate) mod types;
 pub(crate) mod unnecessary_type_conversion;

pyrefly/lib/state/lsp/quick_fixes/remove_unused_variable.rs

@@ -0,0 +1,167 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use dupe::Dupe;
+use pyrefly_python::ast::Ast;
+use pyrefly_python::module::Module;
+use pyrefly_util::visit::Visit;
+use ruff_python_ast::ExceptHandler;
+use ruff_python_ast::Expr;
+use ruff_python_ast::ModModule;
+use ruff_python_ast::Stmt;
+use ruff_text_size::Ranged;
+use ruff_text_size::TextRange;
+
+use crate::ModuleInfo;
+use crate::state::lsp::quick_fixes::extract_shared::line_indent_and_start;
+use crate::state::lsp::quick_fixes::extract_shared::needs_pass_after_removal;
+use crate::state::lsp::quick_fixes::extract_shared::statement_removal_range;
+
+/// Removes an isolated assignment only when its value and binding can be discarded.
+pub(crate) fn remove_unused_variable_code_action(
+    module_info: &ModuleInfo,
+    ast: &ModModule,
+    unused_name: &str,
+    unused_range: TextRange,
+) -> Option<(String, Module, TextRange, String)> {
+    if Ast::is_intentionally_unused(unused_name) {
+        return None;
+    }
+    let context = find_unused_assignment_context(&ast.body, unused_range)?;
+
+    // The unused-variable detector can miss references before an assignment.
+    // Reject other occurrences, including closures and potentially shadowed names.
+    let mut expressions = Vec::new();
+    ast.visit(&mut |expr: &Expr| expressions.push(expr));
+    while let Some(expr) = expressions.pop() {
+        if let Expr::Name(name) = expr
+            && name.id.as_str() == unused_name
+            && name.range() != unused_range
+        {
+            return None;
+        }
+        expr.recurse(&mut |expr| expressions.push(expr));
+    }
+
+    let source = module_info.contents();
+    let removal_range = statement_removal_range(source, context.stmt)?;
+    let prefix = &source[removal_range.start().to_usize()..context.stmt.start().to_usize()];
+    let suffix = &source[context.stmt.end().to_usize()..removal_range.end().to_usize()];
+    let suffix = suffix.trim_start_matches([' ', '\t', '\r', '\n']);
+    // Full-line deletion must not remove a suite header or another statement.
+    // A preceding explicit continuation can also put this statement on a shared logical line.
+    if !prefix.chars().all(|c| matches!(c, ' ' | '\t'))
+        || !(suffix.is_empty() || suffix.starts_with('#'))
+        || source[..removal_range.start().to_usize()]
+            .trim_end_matches(['\r', '\n'])
+            .ends_with('\\')
+    {
+        return None;
+    }
+    let replacement = if needs_pass_after_removal(context.parent_body, context.stmt.range()) {
+        let (indent, _) =
+            line_indent_and_start(module_info.contents(), context.stmt.range().start())?;
+        format!("{indent}pass\n")
+    } else {
+        String::new()
+    };
+    Some((
+        format!("Remove unused variable `{unused_name}`"),
+        module_info.dupe(),
+        removal_range,
+        replacement,
+    ))
+}
+
+struct UnusedAssignmentContext<'a> {
+    stmt: &'a Stmt,
+    parent_body: &'a [Stmt],
+}
+
+fn find_unused_assignment_context(
+    body: &[Stmt],
+    unused_range: TextRange,
+) -> Option<UnusedAssignmentContext<'_>> {
+    for stmt in body {
+        if removable_assignment_range(stmt) == Some(unused_range) {
+            return Some(UnusedAssignmentContext {
+                stmt,
+                parent_body: body,
+            });
+        }
+        if let Some(context) = find_unused_assignment_context_in_child_body(stmt, unused_range) {
+            return Some(context);
+        }
+    }
+    None
+}
+
+fn find_unused_assignment_context_in_child_body(
+    stmt: &Stmt,
+    unused_range: TextRange,
+) -> Option<UnusedAssignmentContext<'_>> {
+    match stmt {
+        Stmt::FunctionDef(function_def) => {
+            find_unused_assignment_context(&function_def.body, unused_range)
+        }
+        Stmt::ClassDef(class_def) => find_unused_assignment_context(&class_def.body, unused_range),
+        Stmt::For(for_stmt) => find_unused_assignment_context(&for_stmt.body, unused_range)
+            .or_else(|| find_unused_assignment_context(&for_stmt.orelse, unused_range)),
+        Stmt::While(while_stmt) => find_unused_assignment_context(&while_stmt.body, unused_range)
+            .or_else(|| find_unused_assignment_context(&while_stmt.orelse, unused_range)),
+        Stmt::If(if_stmt) => {
+            find_unused_assignment_context(&if_stmt.body, unused_range).or_else(|| {
+                if_stmt
+                    .elif_else_clauses
+                    .iter()
+                    .find_map(|clause| find_unused_assignment_context(&clause.body, unused_range))
+            })
+        }
+        Stmt::With(with_stmt) => find_unused_assignment_context(&with_stmt.body, unused_range),
+        Stmt::Match(match_stmt) => match_stmt
+            .cases
+            .iter()
+            .find_map(|case| find_unused_assignment_context(&case.body, unused_range)),
+        Stmt::Try(try_stmt) => find_unused_assignment_context(&try_stmt.body, unused_range)
+            .or_else(|| {
+                try_stmt.handlers.iter().find_map(|handler| {
+                    let ExceptHandler::ExceptHandler(handler) = handler;
+                    find_unused_assignment_context(&handler.body, unused_range)
+                })
+            })
+            .or_else(|| find_unused_assignment_context(&try_stmt.orelse, unused_range))
+            .or_else(|| find_unused_assignment_context(&try_stmt.finalbody, unused_range)),
+        _ => None,
+    }
+}
+
+fn removable_assignment_range(stmt: &Stmt) -> Option<TextRange> {
+    let (target, value) = match stmt {
+        Stmt::Assign(assign) if assign.targets.len() == 1 => {
+            (&assign.targets[0], assign.value.as_ref())
+        }
+        Stmt::AnnAssign(assign) => (assign.target.as_ref(), assign.value.as_deref()?),
+        _ => return None,
+    };
+    // Only literal values are known to have no observable evaluation here.
+    // Other expressions can call user code, raise, suspend, or bind another name.
+    if !matches!(
+        value,
+        Expr::NumberLiteral(_)
+            | Expr::StringLiteral(_)
+            | Expr::BytesLiteral(_)
+            | Expr::BooleanLiteral(_)
+            | Expr::NoneLiteral(_)
+            | Expr::EllipsisLiteral(_)
+    ) {
+        return None;
+    }
+    let Expr::Name(name) = target else {
+        return None;
+    };
+    Some(name.range())
+}

pyrefly/lib/test/lsp/code_actions.rs

@@ -2885,6 +2885,132 @@ fn redundant_cast_action_after(code: &str, cursor_offset: usize) -> Option<Strin
     Some(after)
 }
 
+fn remove_unused_variable_action(code: &str, cursor_offset: usize) -> Option<(String, String)> {
+    let (handles, state) = mk_multi_file_state(&[("main", code)], Require::Exports, false);
+    let handle = handles.get("main")?;
+    let transaction = state.transaction();
+    let module_info = transaction.get_module_info(handle)?;
+    let position = TextSize::try_from(cursor_offset).ok()?;
+    let actions = transaction
+        .local_quickfix_code_actions_sorted(
+            handle,
+            TextRange::new(position, position),
+            ImportFormat::Absolute,
+            None,
+        )
+        .unwrap_or_default();
+    let (title, edits) = actions
+        .into_iter()
+        .find(|(title, _)| title.starts_with("Remove unused variable `"))?;
+    let (module, range, patch) = edits.into_iter().next()?;
+    if module.path() != module_info.path() {
+        return None;
+    }
+    let (_before, after) = apply_patch(&module_info, range, patch);
+    Some((title, after))
+}
+
+#[test]
+fn remove_unused_variable_quickfix() {
+    let code = "def f() -> int:\n    unused = 1\n    return 2\n";
+    let (title, after) = remove_unused_variable_action(code, code.find("unused").unwrap()).unwrap();
+    assert_eq!("Remove unused variable `unused`", title);
+    assert_eq!("def f() -> int:\n    return 2\n", after);
+}
+
+#[test]
+fn remove_unused_variable_inserts_pass_for_empty_body() {
+    let code = "def f() -> None:\n    unused = 1\n";
+    let (_title, after) =
+        remove_unused_variable_action(code, code.find("unused").unwrap()).unwrap();
+    assert_eq!("def f() -> None:\n    pass\n", after);
+}
+
+#[test]
+fn remove_unused_variable_skips_unpacking_assignment() {
+    let code = "def f(pair):\n    unused, used = pair\n    return used\n";
+    assert!(remove_unused_variable_action(code, code.find("unused").unwrap()).is_none());
+}
+
+#[test]
+fn remove_unused_variable_keeps_required_evaluation() {
+    for code in [
+        "def f():\n    unused = print('important')\n",
+        "def f():\n    unused = yield 1\n",
+        "def f():\n    unused = yield from [1]\n",
+        "async def f():\n    unused = await compute()\n",
+        "def f() -> int:\n    unused = (used := 1)\n    return used\n",
+        "def f():\n    unused = 1 / 0\n",
+        "def f():\n    unused = f'{print(1)}'\n",
+        "def f():\n    unused: int = print('important')\n",
+    ] {
+        assert!(
+            remove_unused_variable_action(code, code.find("unused").unwrap()).is_none(),
+            "Unexpected deletion for {code:?}"
+        );
+    }
+}
+
+#[test]
+fn remove_unused_variable_skips_shared_lines() {
+    for code in [
+        "def f() -> int:\n    unused = 1; return 2\n",
+        "def f() -> int:\n    print(1); unused = 1\n    return 2\n",
+        "def f() -> None: unused = 1\n",
+        "def f():\n    unused = 1;\n",
+        "def f():\n    print(1); \\\n        unused = 1\n",
+    ] {
+        assert!(
+            remove_unused_variable_action(code, code.find("unused").unwrap()).is_none(),
+            "Unexpected deletion for {code:?}"
+        );
+    }
+}
+
+#[test]
+fn remove_unused_variable_keeps_referenced_names() {
+    for code in [
+        "def f():\n    def g(): return unused\n    unused = 1\n    return g\n",
+        "def f():\n    for i in range(2):\n        if i:\n            print(unused)\n        unused = 1\n",
+        "def f():\n    g = lambda: unused\n    unused = 1\n    return g\n",
+        "def f():\n    del unused\n    unused = 1\n",
+    ] {
+        let cursor = code.rfind("unused =").unwrap();
+        assert!(
+            remove_unused_variable_action(code, cursor).is_none(),
+            "Unexpected deletion for {code:?}"
+        );
+    }
+}
+
+#[test]
+fn remove_unused_variable_skips_intentionally_unused_names() {
+    for name in ["_", "_unused", "__tracebackhide__"] {
+        let code = format!("def f():\n    {name} = 1\n");
+        assert!(remove_unused_variable_action(&code, code.find(name).unwrap()).is_none());
+    }
+}
+
+#[test]
+fn remove_unused_variable_standalone_lines() {
+    for (code, expected) in [
+        ("def f():\n    unused: int = 1\n", "def f():\n    pass\n"),
+        (
+            "def f():\n    unused = (\n        1\n    )\n",
+            "def f():\n    pass\n",
+        ),
+        (
+            "def f():\n    unused = 1 # comment\n",
+            "def f():\n    pass\n",
+        ),
+        ("def f():\n    unused = 1", "def f():\n    pass\n"),
+        ("def f():\r\n    unused = 1\r\n", "def f():\r\n    pass\r\n"),
+    ] {
+        let (_, after) = remove_unused_variable_action(code, code.find("unused").unwrap()).unwrap();
+        assert_eq!(expected, after);
+    }
+}
+
 #[test]
 fn redundant_cast_parenthesized_expr() {
     let code = "from typing import cast\na: int = 1\nb: int = 2\ncast(int, a + b)\n";