fix code action for Call to deprecated overload contextlib.contextmanager #4245

pyrefly/lib/alt/overload.rs

@@ -44,13 +44,15 @@ use crate::alt::unwrap::HintRef;
 use crate::config::error_kind::ErrorKind;
 use crate::error::collector::ErrorCollector;
 use crate::error::context::ErrorContext;
+use crate::error::error::ErrorQuickFix;
 use crate::solver::solver::ArgumentKey;
 use crate::solver::solver::OverloadTable;
 use crate::solver::solver::TypeVarSpecializationError;
 use crate::types::callable::Callable;
 use crate::types::callable::Params;
 use crate::types::function::FuncMetadata;
 use crate::types::function::Function;
+use crate::types::function::FunctionKind;
 use crate::types::literal::Lit;
 use crate::types::quantified::Quantified;
 use crate::types::type_var::Restriction;
@@ -516,6 +518,23 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 if let Some(detail) = detail {
                     error_builder = error_builder.with_detail(detail);
                 }
+                if let FunctionKind::Def(id) = &closest_overload.func.1.metadata.kind {
+                    let replacement = if id.has_toplevel_qname("contextlib", "contextmanager") {
+                        Some(("Iterator", "Generator"))
+                    } else if id.has_toplevel_qname("contextlib", "asynccontextmanager") {
+                        Some(("AsyncIterator", "AsyncGenerator"))
+                    } else {
+                        None
+                    };
+                    if let Some((from, to)) = replacement {
+                        error_builder = error_builder.with_quick_fix(
+                            ErrorQuickFix::ReplaceDeprecatedContextManagerReturn {
+                                from: from.to_owned(),
+                                to: to.to_owned(),
+                            },
+                        );
+                    }
+                }
                 error_builder.with_context(context).emit();
             }
             errors.extend(closest_overload.arg_errors);

pyrefly/lib/error/error.rs

@@ -58,6 +58,7 @@ pub enum ErrorQuickFix {
     ReplaceWithEnumMember { replacement: String },
     AssertNotNone,
     RemoveUnusedSuppression(SuppressionEdit),
+    ReplaceDeprecatedContextManagerReturn { from: String, to: String },
 }
 
 impl From<SuppressionEdit> for ErrorQuickFix {

pyrefly/lib/error/suppress.rs

@@ -139,7 +139,9 @@ impl SerializedError {
         if let ModulePathDetails::FileSystem(path) = error.path().details() {
             let removal_edit = error.quick_fixes().iter().find_map(|fix| match fix {
                 ErrorQuickFix::RemoveUnusedSuppression(edit) => Some(edit),
-                ErrorQuickFix::ReplaceWithEnumMember { .. } | ErrorQuickFix::AssertNotNone => None,
+                ErrorQuickFix::ReplaceWithEnumMember { .. }
+                | ErrorQuickFix::AssertNotNone
+                | ErrorQuickFix::ReplaceDeprecatedContextManagerReturn { .. } => None,
             });
             let line = removal_edit
                 .map(|edit| {

pyrefly/lib/state/lsp.rs

@@ -3638,6 +3638,18 @@ impl<'a> Transaction<'a> {
                     other_actions.push(action);
                 }
             }
+            if error_range.contains_range(range)
+                && let Some(action) =
+                    quick_fixes::deprecated_contextmanager::replace_return_code_action(
+                        self,
+                        handle,
+                        &module_info,
+                        &ast,
+                        &error,
+                    )
+            {
+                multi_actions.push(action);
+            }
             if error_range.contains_range(range)
                 && let Some(action) = quick_fixes::pyrefly_ignore::add_pyrefly_ignore_code_action(
                     &module_info,
@@ -3783,10 +3795,11 @@ impl<'a> Transaction<'a> {
                         let mut edits = vec![(module, decorator_range, insert_text)];
                         // Import `typing.override` if necessary.
                         if !quick_fixes::add_override::override_in_scope(ast.as_ref())
-                            && let Some(import_edit) = self.override_import_edit(
+                            && let Some(import_edit) = self.typing_import_edit(
                                 handle,
                                 &module_info,
                                 &ast,
+                                "override",
                                 import_format,
                                 custom_thread_pool,
                             )
@@ -3838,19 +3851,19 @@ impl<'a> Transaction<'a> {
         (!actions.is_empty()).then_some(actions)
     }
 
-    /// Builds an edit inserting `from typing import override` (preferring `typing`
-    /// over `typing_extensions`) at the top of the file. Returns `None` when no
-    /// module in scope exports `override`.
-    fn override_import_edit(
+    /// Builds an edit importing `export_name`, preferring `typing` when it exports
+    /// the name. Returns `None` when the name cannot be found.
+    fn typing_import_edit(
         &self,
         handle: &Handle,
         module_info: &Module,
         ast: &ModModule,
+        export_name: &str,
         import_format: ImportFormat,
         custom_thread_pool: Option<&ThreadPool>,
     ) -> Option<(Module, TextRange, String)> {
         let handle_to_import_from = self
-            .search_exports_exact("override", custom_thread_pool)
+            .search_exports_exact(export_name, custom_thread_pool)
             .unwrap_or_default()
             .into_iter()
             .map(|(handle_to_import_from, _, _)| handle_to_import_from)
@@ -3860,7 +3873,7 @@ impl<'a> Transaction<'a> {
             self.config_finder(),
             handle.dupe(),
             handle_to_import_from,
-            "override",
+            export_name,
             import_format,
         );
         Some((module_info.dupe(), edit.range, edit.insert_text))

pyrefly/lib/state/lsp/quick_fixes.rs

@@ -10,6 +10,7 @@ pub(crate) mod assert_not_none;
 pub(crate) mod change_signature;
 pub(crate) mod convert_dict;
 pub(crate) mod convert_star_import;
+pub(crate) mod deprecated_contextmanager;
 pub(crate) mod enum_member;
 pub(crate) mod extract_field;
 pub(crate) mod extract_function;

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

@@ -0,0 +1,126 @@
+/*
+ * 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_build::handle::Handle;
+use pyrefly_python::module::Module;
+use pyrefly_python::sys_info::PythonVersion;
+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 ruff_text_size::TextSize;
+
+use crate::ModuleInfo;
+use crate::error::error::Error;
+use crate::error::error::ErrorQuickFix;
+use crate::state::ide::prepare_import_insertion;
+use crate::state::lsp::Transaction;
+use crate::state::lsp::quick_fixes::extract_shared::find_enclosing_function;
+use crate::state::lsp::quick_fixes::extract_shared::unique_name;
+use crate::types::types::Type;
+
+/// Replaces the deprecated iterator return annotation on a context manager.
+///
+/// Only direct iterator annotations on the diagnosed decorator are editable.
+pub(crate) fn replace_return_code_action(
+    transaction: &Transaction<'_>,
+    handle: &Handle,
+    module_info: &ModuleInfo,
+    ast: &ModModule,
+    error: &Error,
+) -> Option<(String, Vec<(Module, TextRange, String)>)> {
+    let (from, to) = error.quick_fixes().iter().find_map(|fix| match fix {
+        ErrorQuickFix::ReplaceDeprecatedContextManagerReturn { from, to } => {
+            Some((from.as_str(), to.as_str()))
+        }
+        _ => None,
+    })?;
+    let function_def = find_enclosing_function(ast, error.range())?;
+    if !function_def
+        .decorator_list
+        .iter()
+        .any(|decorator| decorator.range() == error.range())
+    {
+        return None;
+    }
+    let annotation = function_def.returns.as_deref()?;
+    let annotation_base = match annotation {
+        Expr::Subscript(subscript) => subscript.value.as_ref(),
+        annotation_base => annotation_base,
+    };
+    let position = match annotation_base {
+        Expr::Name(name) => name.range().start(),
+        Expr::Attribute(attribute) => attribute.attr.range().start(),
+        _ => return None,
+    };
+    // A bound alias can already contain type arguments. Its outer name is not
+    // interchangeable with the iterator class, even if its value is an iterator.
+    let Type::ClassDef(class) = transaction.get_type_at(handle, position)? else {
+        return None;
+    };
+    if !class.has_toplevel_qname("typing", from)
+        && !class.has_toplevel_qname("collections.abc", from)
+    {
+        return None;
+    }
+
+    // These standard modules export both iterator and generator classes. Other
+    // modules can re-export just the iterator, so use a fresh typing import.
+    let qualifier = if let Expr::Attribute(attribute) = annotation_base
+        && let Some(Type::Module(module)) =
+            transaction.get_type_at(handle, attribute.value.range().end() - TextSize::new(1))
+        && matches!(module.to_string().as_str(), "typing" | "collections.abc")
+    {
+        Some(module_info.code_at(attribute.value.range()))
+    } else {
+        None
+    };
+    let mut edits = Vec::new();
+    let binding = if let Some(qualifier) = qualifier {
+        format!("{qualifier}.{to}")
+    } else {
+        // Wildcard imports can bind names that do not occur in the source.
+        if ast.body.iter().any(|stmt| {
+            matches!(stmt, Stmt::ImportFrom(import) if import.names.iter().any(|alias| alias.name.id == "*"))
+        }) {
+            return None;
+        }
+        let binding = unique_name(to, |name| module_info.contents().contains(name));
+        let import = if binding == to {
+            format!("from typing import {to}\n")
+        } else {
+            format!("from typing import {to} as {binding}\n")
+        };
+        let (position, import) = prepare_import_insertion(ast, import);
+        edits.push((module_info.dupe(), TextRange::empty(position), import));
+        binding
+    };
+    let replacement = if let Expr::Subscript(subscript) = annotation {
+        let argument = match subscript.slice.as_ref() {
+            Expr::Tuple(tuple) if tuple.elts.len() == 1 => &tuple.elts[0],
+            Expr::Tuple(_) => return None,
+            argument => argument,
+        };
+        let argument = module_info.code_at(argument.range());
+        let defaults = if handle.sys_info().version() < PythonVersion::new(3, 13, 0) {
+            if to == "Generator" {
+                ", None, None"
+            } else {
+                ", None"
+            }
+        } else {
+            ""
+        };
+        format!("{binding}[{argument}{defaults}]")
+    } else {
+        binding
+    };
+    edits.push((module_info.dupe(), annotation.range(), replacement));
+    Some((format!("Replace `{from}` with `{to}`"), edits))
+}

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

@@ -35,7 +35,9 @@ pub(crate) fn replace_with_enum_member_code_action(
 fn enum_member_replacement(error: &Error) -> Option<&str> {
     error.quick_fixes().iter().find_map(|fix| match fix {
         ErrorQuickFix::ReplaceWithEnumMember { replacement } => Some(replacement.as_str()),
-        ErrorQuickFix::AssertNotNone | ErrorQuickFix::RemoveUnusedSuppression(_) => None,
+        ErrorQuickFix::AssertNotNone
+        | ErrorQuickFix::RemoveUnusedSuppression(_)
+        | ErrorQuickFix::ReplaceDeprecatedContextManagerReturn { .. } => None,
     })
 }
 

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

@@ -25,7 +25,9 @@ pub(crate) fn remove_unused_ignore_code_action(
 
     let edit = error.quick_fixes().iter().find_map(|fix| match fix {
         ErrorQuickFix::RemoveUnusedSuppression(edit) => Some(edit),
-        ErrorQuickFix::ReplaceWithEnumMember { .. } | ErrorQuickFix::AssertNotNone => None,
+        ErrorQuickFix::ReplaceWithEnumMember { .. }
+        | ErrorQuickFix::AssertNotNone
+        | ErrorQuickFix::ReplaceDeprecatedContextManagerReturn { .. } => None,
     })?;
     let title = if edit.replacement.is_empty() {
         "Remove suppression comment"