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

pyrefly/lib/alt/overload.rs

@@ -11,6 +11,7 @@ use itertools::Either;
 use itertools::Itertools;
 use pyrefly_types::callable::ArgCount;
 use pyrefly_types::callable::ArgCounts;
+use pyrefly_types::callable::FunctionKind;
 use pyrefly_types::callable::Param;
 use pyrefly_types::callable::ParamOverlay;
 use pyrefly_types::display::TypeDisplayContext;
@@ -44,6 +45,7 @@ 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::TypeVarSpecializationError;
 use crate::types::callable::Callable;
 use crate::types::callable::FuncMetadata;
@@ -429,6 +431,23 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 if let Some(detail) = detail {
                     error_builder = error_builder.with_detail(detail);
                 }
+                if let FunctionKind::Def(id) = &closest_overload.func.1.metadata.kind
+                    && id.module.name().as_str() == "contextlib"
+                {
+                    let replacement = match id.name.as_str() {
+                        "contextmanager" => Some(("Iterator", "Generator")),
+                        "asynccontextmanager" => Some(("AsyncIterator", "AsyncGenerator")),
+                        _ => 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

@@ -47,6 +47,7 @@ pub struct SecondaryAnnotation {
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum ErrorQuickFix {
     ReplaceWithEnumMember { replacement: String },
+    ReplaceDeprecatedContextManagerReturn { from: String, to: String },
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]

pyrefly/lib/state/lsp.rs

@@ -3018,6 +3018,29 @@ impl<'a> Transaction<'a> {
                     other_actions.push(action);
                 }
             }
+            if error_range.contains_range(range)
+                && let Some((title, module, edit_range, replacement, needs_import)) =
+                    quick_fixes::deprecated_contextmanager::replace_return_code_action(
+                        &module_info,
+                        &ast,
+                        &error,
+                    )
+            {
+                let mut edits = vec![(module, edit_range, replacement.clone())];
+                if needs_import
+                    && let Some(import_edit) = self.typing_import_edit(
+                        handle,
+                        &module_info,
+                        &ast,
+                        &replacement,
+                        import_format,
+                        custom_thread_pool,
+                    )
+                {
+                    edits.push(import_edit);
+                }
+                multi_actions.push((title, edits));
+            }
             if error_range.contains_range(range)
                 && let Some(action) = quick_fixes::pyrefly_ignore::add_pyrefly_ignore_code_action(
                     &module_info,
@@ -3142,10 +3165,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,
                             )
@@ -3197,19 +3221,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)
@@ -3219,7 +3243,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/deprecated_contextmanager.rs

@@ -0,0 +1,78 @@
+/*
+ * 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::module::Module;
+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::error::error::Error;
+use crate::error::error::ErrorQuickFix;
+use crate::state::lsp::quick_fixes::extract_shared::find_enclosing_function;
+
+/// Replaces the deprecated iterator return annotation on a context manager.
+///
+/// The diagnostic identifies the semantic replacement, while the AST locates
+/// the spelling used by the decorated function (including imported aliases).
+pub(crate) fn replace_return_code_action(
+    module_info: &ModuleInfo,
+    ast: &ModModule,
+    error: &Error,
+) -> Option<(String, Module, TextRange, String, bool)> {
+    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())?;
+    let annotation = function_def.returns.as_deref()?;
+    let annotation_base = match annotation {
+        Expr::Subscript(subscript) => subscript.value.as_ref(),
+        annotation_base => annotation_base,
+    };
+    let (range, needs_import) = match annotation_base {
+        Expr::Name(name) => (name.range(), true),
+        Expr::Attribute(attribute) => (attribute.attr.range(), false),
+        _ => return None,
+    };
+    Some((
+        format!("Replace `{from}` with `{to}`"),
+        module_info.dupe(),
+        range,
+        to.to_owned(),
+        needs_import && !bare_name_in_scope(ast, to),
+    ))
+}
+
+fn bare_name_in_scope(ast: &ModModule, name: &str) -> bool {
+    ast.body.iter().any(|stmt| {
+        let Stmt::ImportFrom(import_from) = stmt else {
+            return false;
+        };
+        if import_from.level != 0
+            || !import_from
+                .module
+                .as_ref()
+                .is_some_and(|module| matches!(module.id.as_str(), "typing" | "collections.abc"))
+        {
+            return false;
+        }
+        import_from.names.iter().any(|alias| {
+            alias.name.id.as_str() == "*"
+                || (alias.name.id.as_str() == name
+                    && alias
+                        .asname
+                        .as_ref()
+                        .is_none_or(|asname| asname.id.as_str() == name))
+        })
+    })
+}

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

@@ -33,8 +33,10 @@ pub(crate) fn replace_with_enum_member_code_action(
 }
 
 fn enum_member_replacement(error: &Error) -> Option<&str> {
-    let ErrorQuickFix::ReplaceWithEnumMember { replacement } = error.quick_fixes().first()?;
-    Some(replacement.as_str())
+    error.quick_fixes().iter().find_map(|fix| match fix {
+        ErrorQuickFix::ReplaceWithEnumMember { replacement } => Some(replacement.as_str()),
+        _ => None,
+    })
 }
 
 fn enclosing_string_literal_range(ast: &ModModule, error_range: TextRange) -> Option<TextRange> {

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

@@ -8,6 +8,7 @@
 pub(crate) mod add_override;
 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/test/lsp/code_actions.rs

@@ -1031,6 +1031,84 @@ takes_status("active")
     );
 }
 
+#[test]
+fn quickfix_deprecated_contextmanager_return() {
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[(
+            "main",
+            r#"from contextlib import contextmanager
+from typing import Iterator
+
+@contextmanager
+# ^
+def managed() -> Iterator[int]:
+    yield 1
+"#,
+        )],
+        get_test_report,
+    );
+    assert!(
+        report.contains("# Title: Replace `Iterator` with `Generator`"),
+        "{report}"
+    );
+    assert!(report.contains("from typing import Generator"), "{report}");
+    assert!(
+        report.contains("def managed() -> Generator[int]:"),
+        "{report}"
+    );
+}
+
+#[test]
+fn quickfix_deprecated_contextmanager_qualified_and_async_returns() {
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[
+            (
+                "sync",
+                r#"import typing
+from contextlib import contextmanager
+
+@contextmanager
+# ^
+def managed() -> typing.Iterator[int]:
+    yield 1
+"#,
+            ),
+            (
+                "async",
+                r#"from contextlib import asynccontextmanager
+from typing import AsyncIterator
+
+@asynccontextmanager
+# ^
+async def managed() -> AsyncIterator[int]:
+    yield 1
+"#,
+            ),
+        ],
+        get_test_report,
+    );
+    assert!(
+        report.contains("def managed() -> typing.Generator[int]:"),
+        "{report}"
+    );
+    assert!(
+        !report.contains("from typing import Generator\n"),
+        "{report}"
+    );
+    assert!(
+        report.contains("# Title: Replace `AsyncIterator` with `AsyncGenerator`"),
+        "{report}"
+    );
+    assert!(
+        report.contains("from typing import AsyncGenerator"),
+        "{report}"
+    );
+    assert!(
+        report.contains("async def managed() -> AsyncGenerator[int]:"),
+        "{report}"
+    );
+}
+
 #[test]
 fn quickfix_add_pyrefly_ignore_code_with_existing_comment() {
     let report = get_batched_lsp_operations_report_allow_error(