From f3eeca24a8fcab6694bf9cc5b05cedfadf84f683 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Fri, 30 Oct 2026 12:49:43 +0900 Subject: [PATCH 1/2] test --- pyrefly/lib/test/lsp/code_actions.rs | 312 +++++++++++++++++++++++++++ 1 file changed, 312 insertions(+) diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 2b83c8698a..6cd57f51bc 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -12,6 +12,7 @@ use pyrefly_build::handle::Handle; use pyrefly_python::ast::Ast; use pyrefly_python::module::Module; use ruff_python_ast::PySourceType; +use pyrefly_python::sys_info::PythonVersion; use ruff_text_size::TextRange; use ruff_text_size::TextSize; @@ -29,6 +30,7 @@ use crate::test::util::extract_cursors_for_test; use crate::test::util::get_batched_lsp_operations_report_allow_error; use crate::test::util::mk_multi_file_state; use crate::test::util::mk_multi_file_state_assert_no_errors; +use crate::test::util::mk_multi_file_state_with_env; fn apply_patch(info: &ModuleInfo, range: TextRange, patch: String) -> (String, String) { let before = info.contents().as_str().to_owned(); @@ -1113,6 +1115,31 @@ fn insert_import_uses_file_line_ending() { ); } +#[test] +fn insert_import_after_future_imports() { + for (unknown, import) in [ + ("my_export", "from a import my_export"), + ("np", "import numpy as np"), + ] { + let code = format!( + "\"\"\"Module documentation.\"\"\"\n\ + from __future__ import annotations\n\ + from __future__ import generator_stop\n\ + {unknown}\n# ^\n" + ); + let report = get_batched_lsp_operations_report_allow_error( + &[("a", "my_export = 3\n"), ("numpy", ""), ("main", &code)], + get_test_report, + ); + assert!( + report.contains(&format!( + "from __future__ import generator_stop\n{import}\n{unknown}" + )), + "{report}" + ); + } +} + #[test] fn insertion_test_comments() { let report = get_batched_lsp_operations_report_allow_error( @@ -1391,6 +1418,291 @@ def f(x: T | None) -> None: ); } +#[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_deprecated_contextmanager_python_versions() { + for minor in [10, 11, 12, 13] { + for (decorator, iterator, generator, async_prefix, defaults) in [ + ( + "contextmanager", + "Iterator", + "Generator", + "", + ", None, None", + ), + ( + "asynccontextmanager", + "AsyncIterator", + "AsyncGenerator", + "async ", + ", None", + ), + ] { + let code = format!( + "from contextlib import {decorator}\nfrom typing import {iterator}\n\n\ + @{decorator}\n# ^\n{async_prefix}def managed() -> {iterator}[int]:\n yield 1\n" + ); + let (handles, state) = mk_multi_file_state_with_env( + TestEnv::new_with_version(PythonVersion::new(3, minor, 0)), + &[("main", &code)], + Require::Exports, + false, + ); + let report = + get_test_report(&state, &handles["main"], extract_cursors_for_test(&code)[0]); + let defaults = if minor < 13 { defaults } else { "" }; + assert!( + report.contains(&format!( + "{async_prefix}def managed() -> {generator}[int{defaults}]:" + )), + "Python 3.{minor}: {report}" + ); + } + } +} + +#[test] +fn quickfix_deprecated_contextmanager_preserves_future_imports() { + for line_ending in ["\n", "\r\n"] { + let code = r#""""Module documentation.""" +from __future__ import annotations +from __future__ import generator_stop +from contextlib import contextmanager +from typing import Iterator + +@contextmanager +# ^ +def managed() -> Iterator[int]: + yield 1 +"# + .replace('\n', line_ending); + let report = + get_batched_lsp_operations_report_allow_error(&[("main", &code)], get_test_report); + assert!( + report.contains(&format!( + "from __future__ import generator_stop{line_ending}\ + from typing import Generator{line_ending}\ + from contextlib import contextmanager" + )), + "{report}" + ); + assert!( + report.contains("def managed() -> Generator[int]:"), + "{report}" + ); + } +} + +#[test] +fn quickfix_deprecated_contextmanager_explicit_call_not_offered() { + let report = get_batched_lsp_operations_report_allow_error( + &[( + "main", + r#"from contextlib import contextmanager +from typing import Iterator + +def producer() -> Iterator[int]: + yield 1 + +def wrapper() -> int: + contextmanager(producer) +# ^ + return 1 +"#, + )], + get_test_report_with_error_count, + ); + assert!(report.contains("Error count: 1"), "{report}"); + assert!(!report.contains("# Title: Replace `Iterator`"), "{report}"); +} + +#[test] +fn quickfix_deprecated_contextmanager_wrappers_and_bound_aliases_not_offered() { + for annotation in [ + "Annotated[Iterator[int], 'tag']", + "Bound", + "BoundAlias", + "Alias[int]", + ] { + let code = format!( + "from contextlib import contextmanager\n\ + from typing import Annotated, Iterator, TypeAlias, TypeVar\n\ + T = TypeVar('T')\n\ + Bound = Iterator[int]\n\ + BoundAlias: TypeAlias = Iterator[int]\n\ + Alias = Iterator[T]\n\n\ + @contextmanager\n# ^\ndef managed() -> {annotation}:\n yield 1\n" + ); + let report = get_batched_lsp_operations_report_allow_error( + &[("main", &code)], + get_test_report_with_error_count, + ); + assert!(report.contains("Error count: 1"), "{report}"); + assert!(!report.contains("# Title: Replace `Iterator`"), "{report}"); + } +} + +#[test] +fn quickfix_deprecated_contextmanager_reexport_and_shadowed_binding() { + for (annotation, imports, binding) in [ + ("api.Iterator[int]", "import api", "Generator"), + ( + "It[int]", + "from typing import Iterator as It\nGenerator = int", + "Generator_1", + ), + ( + "It[int]", + "from typing import Iterator as It, Generator\nGenerator = int", + "Generator_1", + ), + ] { + let code = format!( + "from contextlib import contextmanager\n{imports}\n\n\ + @contextmanager\n# ^\ndef managed() -> {annotation}:\n yield 1\n" + ); + let report = get_batched_lsp_operations_report_allow_error( + &[ + ("api", "from typing import Iterator as Iterator\n"), + ("main", &code), + ], + get_test_report, + ); + assert!( + report.contains(&format!("def managed() -> {binding}[int]:")), + "{report}" + ); + let import = if binding == "Generator" { + "from typing import Generator\n" + } else { + "from typing import Generator as Generator_1\n" + }; + assert!(report.contains(import), "{report}"); + assert!(!report.contains("api.Generator"), "{report}"); + } +} + +#[test] +fn quickfix_deprecated_contextmanager_nested_shadowed_binding() { + let report = get_batched_lsp_operations_report_allow_error( + &[( + "main", + r#"from contextlib import contextmanager +from typing import Iterator + +def wrapper(Generator: type[int]) -> None: + @contextmanager +# ^ + def managed() -> Iterator[int]: + yield 1 +"#, + )], + get_test_report, + ); + assert!( + report.contains("from typing import Generator as Generator_1\n"), + "{report}" + ); + assert!( + report.contains("def managed() -> Generator_1[int]:"), + "{report}" + ); +} + +#[test] +fn quickfix_deprecated_contextmanager_qualified_module_aliases() { + for module in ["typing", "collections.abc"] { + let code = format!( + "import {module} as abc\nfrom contextlib import contextmanager\n\n\ + @contextmanager\n# ^\ndef managed() -> abc.Iterator[tuple[int, str]]:\n yield (1, 'a')\n" + ); + let report = + get_batched_lsp_operations_report_allow_error(&[("main", &code)], get_test_report); + assert!( + report.contains("def managed() -> abc.Generator[tuple[int, str]]:"), + "{report}" + ); + assert!(!report.contains("from typing import Generator"), "{report}"); + } +} + #[test] fn quickfix_add_pyrefly_ignore_code_with_existing_comment() { let report = get_batched_lsp_operations_report_allow_error( From 147a1982eafe00f71bf2fea8353c1c512f1fc974 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 25 Jul 2026 01:17:10 +0900 Subject: [PATCH 2/2] fix --- pyrefly/lib/alt/overload.rs | 19 +++ pyrefly/lib/error/error.rs | 1 + pyrefly/lib/error/suppress.rs | 4 +- pyrefly/lib/state/lsp.rs | 27 +++- pyrefly/lib/state/lsp/quick_fixes.rs | 1 + .../quick_fixes/deprecated_contextmanager.rs | 126 ++++++++++++++++++ .../lib/state/lsp/quick_fixes/enum_member.rs | 4 +- .../state/lsp/quick_fixes/unused_ignore.rs | 4 +- pyrefly/lib/test/lsp/code_actions.rs | 2 +- 9 files changed, 177 insertions(+), 11 deletions(-) create mode 100644 pyrefly/lib/state/lsp/quick_fixes/deprecated_contextmanager.rs diff --git a/pyrefly/lib/alt/overload.rs b/pyrefly/lib/alt/overload.rs index 2a12fdb145..937648f783 100644 --- a/pyrefly/lib/alt/overload.rs +++ b/pyrefly/lib/alt/overload.rs @@ -44,6 +44,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::ArgumentKey; use crate::solver::solver::OverloadTable; use crate::solver::solver::TypeVarSpecializationError; @@ -51,6 +52,7 @@ 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); diff --git a/pyrefly/lib/error/error.rs b/pyrefly/lib/error/error.rs index 886b496645..00b36e24e4 100644 --- a/pyrefly/lib/error/error.rs +++ b/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 for ErrorQuickFix { diff --git a/pyrefly/lib/error/suppress.rs b/pyrefly/lib/error/suppress.rs index e1bec1426f..5fb4d482d7 100644 --- a/pyrefly/lib/error/suppress.rs +++ b/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| { diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 8088d344cc..6aba00ed75 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/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)) diff --git a/pyrefly/lib/state/lsp/quick_fixes.rs b/pyrefly/lib/state/lsp/quick_fixes.rs index 2aeed6d263..af551983c7 100644 --- a/pyrefly/lib/state/lsp/quick_fixes.rs +++ b/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; diff --git a/pyrefly/lib/state/lsp/quick_fixes/deprecated_contextmanager.rs b/pyrefly/lib/state/lsp/quick_fixes/deprecated_contextmanager.rs new file mode 100644 index 0000000000..ae2c704a15 --- /dev/null +++ b/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)) +} diff --git a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs index 9da8f37d98..5833a51851 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs +++ b/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, }) } diff --git a/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs b/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs index 7635c08cab..4c1fc41061 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs +++ b/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" diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 6cd57f51bc..f9a046092a 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -11,8 +11,8 @@ use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; use pyrefly_python::ast::Ast; use pyrefly_python::module::Module; -use ruff_python_ast::PySourceType; use pyrefly_python::sys_info::PythonVersion; +use ruff_python_ast::PySourceType; use ruff_text_size::TextRange; use ruff_text_size::TextSize;