From 12f48475ffab646e75fdf53e6bd3111462c0f34d 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 | 588 +++++++++++++++++++++++++++ 1 file changed, 588 insertions(+) diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 3f00528f55..c6ccef7079 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -11,8 +11,10 @@ use pretty_assertions::assert_eq; use pyrefly_build::handle::Handle; use pyrefly_python::ast::Ast; use pyrefly_python::module::Module; +use pyrefly_python::sys_info::PythonVersion; use ruff_python_ast::PySourceType; use ruff_text_size::Ranged; +use pyrefly_python::sys_info::PythonVersion; use ruff_text_size::TextRange; use ruff_text_size::TextSize; @@ -1237,6 +1239,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( @@ -1515,6 +1542,567 @@ 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}" + ); +} + +fn apply_contextmanager_fix(code: &str, minor: u32) -> Option { + let (handles, state) = mk_multi_file_state_with_env( + TestEnv::new_with_version(PythonVersion::new(3, minor, 0)), + &[ + ("api", "Generator = int\nAsyncGenerator = int\n"), + ("main", code), + ], + Require::Exports, + false, + ); + let transaction = state.transaction(); + let handle = &handles["main"]; + for position in extract_cursors_for_test(code) { + let actions = transaction + .local_quickfix_code_actions_sorted( + handle, + TextRange::empty(position), + ImportFormat::Absolute, + None, + ) + .unwrap_or_default(); + if let Some((_, edits)) = actions.iter().find(|(title, _)| { + title == "Replace `Iterator` with `Generator`" + || title == "Replace `AsyncIterator` with `AsyncGenerator`" + }) { + let module = transaction.get_module_info(handle).unwrap(); + let after = apply_refactor_edits_for_module(&module, edits); + let (_, errors, _) = Ast::parse(&after, PySourceType::Python); + assert!(errors.is_empty(), "{errors:?}\n{after}"); + return Some(after); + } + } + None +} + +#[test] +fn quickfix_deprecated_contextmanager_preserves_annotation_source() { + for (decorator, iterator, generator, async_prefix, defaults) in [ + ( + "contextmanager", + "Iterator", + "Generator", + "", + ", None, None", + ), + ( + "asynccontextmanager", + "AsyncIterator", + "AsyncGenerator", + "async ", + ", None", + ), + ] { + for minor in [12, 13] { + for newline in ["\n", "\r\n"] { + for annotation in [ + "ITERATOR[int@@]", + "ITERATOR[int@@,]", + "ITERATOR[\n Literal[\n 'first',\n 'second',\n ]@@, # Keep the trailing comma.\n]", + "ITERATOR[\n ((Literal[\n 'first',\n 'second',\n ]))@@, # Keep the parentheses.\n]", + "(ITERATOR # Keep the name comment.\n)[(Literal['''first\nsecond'''])@@,]", + "ITERATOR[(int # Keep the argument comment.\n)@@,]", + "ITERATOR[(int@@,)]", + ] { + let prefix = format!( + "from contextlib import {decorator}\n\ + from typing import {iterator}, {generator}, Literal, Any\n\n\ + @{decorator}\n# ^\n{async_prefix}def managed(value: Any) -> " + ); + let suffix = ":\n yield value\n"; + let before = format!( + "{prefix}{}{suffix}", + annotation.replace("ITERATOR", iterator).replace("@@", "") + ) + .replace('\n', newline); + let expected = format!( + "{prefix}{}{suffix}", + annotation + .replace("ITERATOR", generator) + .replace("@@", if minor < 13 { defaults } else { "" }) + ) + .replace('\n', newline); + let after = apply_contextmanager_fix(&before, minor) + .unwrap_or_else(|| panic!("No fix for Python 3.{minor}: {before}")); + assert_eq!(after, expected, "Python 3.{minor}"); + assert!(!after.contains("\r\r\n")); + } + } + } + } +} + +#[test] +fn quickfix_deprecated_contextmanager_block_wildcard_imports() { + for block in [ + "if True:\n from api import *\n", + "try:\n from api import *\nexcept ImportError:\n pass\n", + "for item in [1]:\n from api import *\n", + "while False:\n from api import *\n", + ] { + for (decorator, iterator, async_prefix) in [ + ("contextmanager", "Iterator", ""), + ("asynccontextmanager", "AsyncIterator", "async "), + ] { + let code = format!( + "from contextlib import {decorator}\nfrom typing import {iterator}\n\ + {block}\n@{decorator}\n# ^\n{async_prefix}def managed() -> {iterator}[int]:\n yield 1\n" + ); + assert!(apply_contextmanager_fix(&code, 13).is_none(), "{code}"); + } + } +} + +#[test] +fn quickfix_deprecated_contextmanager_normalized_binding_names() { + for (decorator, iterator, generator, normalized_collision, async_prefix) in [ + ("contextmanager", "Iterator", "Generator", "Ｇenerator", ""), + ( + "asynccontextmanager", + "AsyncIterator", + "AsyncGenerator", + "ＡsyncGenerator", + "async ", + ), + ] { + for (scope, indent) in [ + (format!("{normalized_collision} = int\n"), ""), + (format!("def {normalized_collision}():\n pass\n"), ""), + (format!("class {normalized_collision}:\n pass\n"), ""), + (format!("def outer({normalized_collision}: int):\n"), " "), + ( + format!("def outer():\n {normalized_collision} = int\n"), + " ", + ), + ] { + let code = format!( + "from contextlib import {decorator}\nfrom typing import {iterator}\n\ + {scope}{indent}@{decorator}\n#{indent} ^\n\ + {indent}{async_prefix}def managed() -> {iterator}[int]:\n{indent} yield 1\n" + ); + let after = apply_contextmanager_fix(&code, 13).unwrap(); + assert!( + after.contains(&format!( + "from typing import {generator} as {generator}_1\n" + )), + "{after}" + ); + assert!( + after.contains(&format!("def managed() -> {generator}_1[int]:")), + "{after}" + ); + } + } + let code = "from contextlib import contextmanager\n\ + from typing import Iterator, AsyncGenerator\n\ + # Generator and Generator_1 in comments do not bind names.\n\ + GeneratorExit = int\n\n\ + @contextmanager\n# ^\ndef managed() -> Iterator[int]:\n yield 1\n"; + let after = apply_contextmanager_fix(code, 13).unwrap(); + assert!( + after.starts_with("from typing import Generator\n"), + "{after}" + ); + assert!( + after.contains("def managed() -> Generator[int]:"), + "{after}" + ); +} + +#[test] +fn quickfix_deprecated_contextmanager_requires_own_yield() { + for body in [ + " return iter([1])\n", + " def nested():\n yield 1\n return iter([1])\n", + " class Nested:\n def values(self):\n yield 1\n return iter([1])\n", + " nested = lambda: (yield 1)\n return iter([1])\n", + ] { + let code = format!( + "from contextlib import contextmanager\nfrom typing import Iterator\n\n\ + @contextmanager\n# ^\ndef managed() -> Iterator[int]:\n{body}" + ); + assert!(apply_contextmanager_fix(&code, 13).is_none(), "{code}"); + } + let code = "from contextlib import asynccontextmanager\nfrom typing import AsyncIterator\n\ + async def producer() -> AsyncIterator[int]:\n yield 1\n\n\ + @asynccontextmanager\n# ^\ndef managed() -> AsyncIterator[int]:\n return producer()\n"; + assert!(apply_contextmanager_fix(code, 13).is_none()); + let code = "from contextlib import asynccontextmanager\nfrom typing import AsyncIterator\n\n\ + @asynccontextmanager\n# ^\ndef managed() -> AsyncIterator[int]:\n yield 1\n"; + assert!(apply_contextmanager_fix(code, 13).is_none()); + let code = "from contextlib import contextmanager\nfrom typing import Iterator\n\n\ + @contextmanager\n# ^\ndef managed() -> Iterator[int]:\n yield from [1]\n"; + assert!(apply_contextmanager_fix(code, 13).is_some()); +} + +#[test] +fn quickfix_deprecated_contextmanager_reuses_imports_on_repeated_fixes() { + for (decorator, iterator, generator, async_prefix) in [ + ("contextmanager", "Iterator", "Generator", ""), + ( + "asynccontextmanager", + "AsyncIterator", + "AsyncGenerator", + "async ", + ), + ] { + for module in ["typing", "collections.abc"] { + for (import, binding) in [ + ( + format!("from {module} import {iterator}\n"), + generator.to_owned(), + ), + ( + format!("from {module} import {iterator}, {generator}\n"), + generator.to_owned(), + ), + ( + format!("from {module} import {iterator}\n{generator} = int\n"), + format!("{generator}_1"), + ), + ] { + let before = format!( + "from contextlib import {decorator}\n{import}\n\ + @{decorator}\n# ^\n{async_prefix}def first() -> {iterator}[int]:\n yield 1\n\n\ + @{decorator}\n# ^\n{async_prefix}def second() -> {iterator}[int]:\n yield 2\n" + ); + let first = apply_contextmanager_fix(&before, 13).unwrap(); + let second = apply_contextmanager_fix(&first, 13).unwrap(); + assert!( + second.contains(&format!("def first() -> {binding}[int]:")), + "{second}" + ); + assert!( + second.contains(&format!("def second() -> {binding}[int]:")), + "{second}" + ); + assert_eq!( + first.matches(" import ").count(), + second.matches(" import ").count() + ); + assert!(apply_contextmanager_fix(&second, 13).is_none()); + } + } + } +} + +#[test] +fn quickfix_deprecated_contextmanager_does_not_reuse_late_import() { + for (decorator, iterator, generator, async_prefix) in [ + ("contextmanager", "Iterator", "Generator", ""), + ( + "asynccontextmanager", + "AsyncIterator", + "AsyncGenerator", + "async ", + ), + ] { + 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\n\ + from typing import {generator}\n" + ); + let after = apply_contextmanager_fix(&code, 13).unwrap(); + assert!( + after.starts_with(&format!( + "from typing import {generator} as {generator}_1\n" + )), + "{after}" + ); + assert!( + after.contains(&format!("def managed() -> {generator}_1[int]:")), + "{after}" + ); + } +} + +#[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 c34040c135d44af78e2090d53954dcfddd457c69 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 | 17 ++ pyrefly/lib/error/error.rs | 17 ++ pyrefly/lib/error/suppress.rs | 3 +- pyrefly/lib/lsp/non_wasm/server.rs | 6 +- pyrefly/lib/state/lsp.rs | 25 +- pyrefly/lib/state/lsp/quick_fixes.rs | 1 + .../quick_fixes/deprecated_contextmanager.rs | 229 ++++++++++++++++++ .../lib/state/lsp/quick_fixes/enum_member.rs | 3 +- .../state/lsp/quick_fixes/unused_ignore.rs | 3 +- pyrefly/lib/test/lsp/code_actions.rs | 1 - 10 files changed, 290 insertions(+), 15 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..d1280c3f34 100644 --- a/pyrefly/lib/alt/overload.rs +++ b/pyrefly/lib/alt/overload.rs @@ -44,6 +44,8 @@ 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::ContextManagerKind; +use crate::error::error::ErrorQuickFix; use crate::solver::solver::ArgumentKey; use crate::solver::solver::OverloadTable; use crate::solver::solver::TypeVarSpecializationError; @@ -51,6 +53,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 +519,20 @@ 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(ContextManagerKind::Sync) + } else if id.has_toplevel_qname("contextlib", "asynccontextmanager") { + Some(ContextManagerKind::Async) + } else { + None + }; + if let Some(kind) = replacement { + error_builder = error_builder.with_quick_fix( + ErrorQuickFix::ReplaceDeprecatedContextManagerReturn(kind), + ); + } + } 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 083dc7bebb..f8dc296e6a 100644 --- a/pyrefly/lib/error/error.rs +++ b/pyrefly/lib/error/error.rs @@ -60,6 +60,23 @@ pub enum ErrorQuickFix { AssertNotNone, RemoveUnusedSuppression(SuppressionEdit), ImportModule { module: ModuleName }, + ReplaceDeprecatedContextManagerReturn(ContextManagerKind), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ContextManagerKind { + Sync, + Async, +} + +impl ContextManagerKind { + /// The iterator, generator, and explicit defaults for a return annotation. + pub fn return_annotation_replacement(self) -> (&'static str, &'static str, &'static str) { + match self { + Self::Sync => ("Iterator", "Generator", ", None, None"), + Self::Async => ("AsyncIterator", "AsyncGenerator", ", None"), + } + } } impl From for ErrorQuickFix { diff --git a/pyrefly/lib/error/suppress.rs b/pyrefly/lib/error/suppress.rs index bd0ee28b0e..f7343260f1 100644 --- a/pyrefly/lib/error/suppress.rs +++ b/pyrefly/lib/error/suppress.rs @@ -141,7 +141,8 @@ impl SerializedError { ErrorQuickFix::RemoveUnusedSuppression(edit) => Some(edit), ErrorQuickFix::ReplaceWithEnumMember { .. } | ErrorQuickFix::AssertNotNone - | ErrorQuickFix::ImportModule { .. } => None, + | ErrorQuickFix::ImportModule { .. } + | ErrorQuickFix::ReplaceDeprecatedContextManagerReturn(_) => None, }); let line = removal_edit .map(|edit| { diff --git a/pyrefly/lib/lsp/non_wasm/server.rs b/pyrefly/lib/lsp/non_wasm/server.rs index 934958e5fc..4682e87895 100644 --- a/pyrefly/lib/lsp/non_wasm/server.rs +++ b/pyrefly/lib/lsp/non_wasm/server.rs @@ -5247,9 +5247,9 @@ impl Server { Some(&self.lsp_thread_pool), ) { actions.extend(quickfixes.into_iter().filter_map(|(title, edits)| { - // A quick fix may carry more than one edit (e.g. the missing - // `@override` fix inserts both the decorator and an import). Group - // every edit by the document it targets into a single workspace edit + // Fixes for missing `@override` and context-manager annotations + // can include several edits. Group every edit by its target document + // into a single workspace edit // so the whole fix applies in one action. let mut changes: HashMap> = HashMap::new(); for (info, range, insert_text) in edits { diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 1da1f71bfd..0d27d1050e 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -3648,8 +3648,8 @@ impl<'a> Transaction<'a> { let mut import_actions = Vec::new(); let mut generate_actions = Vec::new(); let mut other_actions = Vec::new(); - // Actions that carry more than one edit (e.g. the missing-`@override` fix, - // which inserts both the decorator and an import). + // Actions such as missing-`@override` and context-manager annotations + // carry multiple edits that must apply together. let mut multi_actions: Vec<(String, Vec<(Module, TextRange, String)>)> = Vec::new(); // Deduplicates the actions pushed onto `other_actions`. The same fix can // be generated more than once -- e.g. several errors on one line all @@ -3724,6 +3724,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, @@ -3892,9 +3904,7 @@ impl<'a> Transaction<'a> { // this will be the public/non-deprecated version) import_actions.dedup_by(|a, b| a.insert_text == b.insert_text); - // Every quick-fix producer except the missing-`@override` fix yields a single - // edit; wrap those in a one-element edit list so they share the multi-edit shape - // that `multi_actions` and the LSP layer expect. + // Wrap single-edit actions in the same shape as the multi-edit producers. fn wrap_single( (title, module, range, insert_text): (String, Module, TextRange, String), ) -> (String, Vec<(Module, TextRange, String)>) { @@ -3924,9 +3934,8 @@ 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`. + /// Imports `override`, preferring `typing` over `typing_extensions`. + /// Returns `None` when the name cannot be found. fn override_import_edit( &self, handle: &Handle, diff --git a/pyrefly/lib/state/lsp/quick_fixes.rs b/pyrefly/lib/state/lsp/quick_fixes.rs index f92818f4ec..5bbb87fc63 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..6390f55fc5 --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/deprecated_contextmanager.rs @@ -0,0 +1,229 @@ +/* + * 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 std::collections::HashMap; + +use dupe::Dupe; +use pyrefly_build::handle::Handle; +use pyrefly_python::ast::Ast; +use pyrefly_python::module::Module; +use pyrefly_python::sys_info::PythonVersion; +use ruff_python_ast::Expr; +use ruff_python_ast::ExprContext; +use ruff_python_ast::Identifier; +use ruff_python_ast::ModModule; +use ruff_python_ast::Stmt; +use ruff_python_ast::name::Name; +use ruff_python_ast::token::parenthesized_range; +use ruff_python_ast::visitor::source_order::SourceOrderVisitor; +use ruff_python_ast::visitor::source_order::walk_expr; +use ruff_python_ast::visitor::source_order::walk_stmt; +use ruff_python_parser::parse_expression_range; +use ruff_text_size::Ranged; +use ruff_text_size::TextRange; +use ruff_text_size::TextSize; + +use crate::ModuleInfo; +use crate::error::error::ContextManagerKind; +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 kind = error.quick_fixes().iter().find_map(|fix| match fix { + ErrorQuickFix::ReplaceDeprecatedContextManagerReturn(kind) => Some(*kind), + ErrorQuickFix::ReplaceWithEnumMember { .. } + | ErrorQuickFix::AssertNotNone + | ErrorQuickFix::RemoveUnusedSuppression(_) + | ErrorQuickFix::ImportModule { .. } => None, + })?; + let (from, to, defaults) = kind.return_annotation_replacement(); + let function_def = find_enclosing_function(ast, error.range())?; + if !function_def + .decorator_list + .iter() + .any(|decorator| decorator.range() == error.range()) + || !Ast::body_contains_yield(&function_def.body) + || function_def.is_async != matches!(kind, ContextManagerKind::Async) + { + 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(attribute) + } else { + None + }; + let mut edits = Vec::new(); + let (name_range, binding) = if let Some(attribute) = qualifier { + // Preserve the qualifier and every byte around the iterator name. + (attribute.attr.range(), to.to_owned()) + } else { + let mut names = BindingNames::default(); + names.visit_body(&ast.body); + if names.wildcard_import { + return None; + } + // A sole, earlier standard import cannot be shadowed by another binding. + let existing = ast.body.iter().find_map(|stmt| { + let Stmt::ImportFrom(import) = stmt else { + return None; + }; + if import.end() > function_def.start() + || import.level != 0 + || !import.module.as_ref().is_some_and(|module| { + matches!(module.id.as_str(), "typing" | "collections.abc") + }) + { + return None; + } + import.names.iter().find_map(|alias| { + let binding = alias.asname.as_ref().unwrap_or(&alias.name); + (alias.name.id == to + && names + .sites + .get(&binding.id) + .is_some_and(|sites| sites.as_slice() == [binding.range()])) + .then(|| binding.id.to_string()) + }) + }); + let binding = if let Some(existing) = existing { + existing + } else { + let binding = unique_name(to, |name| names.sites.contains_key(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 + }; + (annotation_base.range(), binding) + }; + edits.push((module_info.dupe(), name_range, binding)); + 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(_) | Expr::Starred(_) => return None, + argument => argument, + }; + if handle.sys_info().version() < PythonVersion::new(3, 13, 0) { + let parsed = parse_expression_range(module_info.contents(), annotation.range()).ok()?; + // Grouping parentheses belong to the original argument. Insert before + // its trailing comma, after all closing parentheses, without copying source. + let range = parenthesized_range(argument.into(), subscript.into(), parsed.tokens()) + .unwrap_or(argument.range()); + edits.push(( + module_info.dupe(), + TextRange::empty(range.end()), + defaults.to_owned(), + )); + } + } + Some((format!("Replace `{from}` with `{to}`"), edits)) +} + +/// Binding sites use the parser's normalized identifiers, across enclosing and +/// nested scopes. Loads, attribute names, and comments do not bind a fresh name. +#[derive(Default)] +struct BindingNames { + sites: HashMap>, + wildcard_import: bool, +} + +impl<'a> SourceOrderVisitor<'a> for BindingNames { + fn visit_identifier(&mut self, identifier: &'a Identifier) { + self.sites + .entry(identifier.id.clone()) + .or_default() + .push(identifier.range()); + } + + fn visit_expr(&mut self, expr: &'a Expr) { + match expr { + Expr::Name(name) => { + if name.ctx != ExprContext::Load { + self.sites + .entry(name.id.clone()) + .or_default() + .push(name.range()); + } + } + Expr::Attribute(attribute) => self.visit_expr(&attribute.value), + _ => walk_expr(self, expr), + } + } + + fn visit_stmt(&mut self, stmt: &'a Stmt) { + match stmt { + Stmt::ImportFrom(import) => { + for alias in &import.names { + self.wildcard_import |= alias.name.id == "*"; + self.visit_identifier(alias.asname.as_ref().unwrap_or(&alias.name)); + } + } + Stmt::Import(import) => { + for alias in &import.names { + if let Some(name) = &alias.asname { + self.visit_identifier(name); + } else { + let path = alias.name.id.as_str(); + let root = path.split_once('.').map_or(path, |(root, _)| root); + self.sites + .entry(Name::new(root)) + .or_default() + .push(alias.name.range()); + } + } + } + // Compound statements remain in the traversal, including module-level + // conditionals, exception handlers, and loops containing wildcard imports. + _ => walk_stmt(self, stmt), + } + } +} diff --git a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs index b17d7be526..12cff70085 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs @@ -37,7 +37,8 @@ fn enum_member_replacement(error: &Error) -> Option<&str> { ErrorQuickFix::ReplaceWithEnumMember { replacement } => Some(replacement.as_str()), ErrorQuickFix::AssertNotNone | ErrorQuickFix::RemoveUnusedSuppression(_) - | ErrorQuickFix::ImportModule { .. } => None, + | ErrorQuickFix::ImportModule { .. } + | 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 d5833a0a21..ca35a907bd 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs @@ -27,7 +27,8 @@ pub(crate) fn remove_unused_ignore_code_action( ErrorQuickFix::RemoveUnusedSuppression(edit) => Some(edit), ErrorQuickFix::ReplaceWithEnumMember { .. } | ErrorQuickFix::AssertNotNone - | ErrorQuickFix::ImportModule { .. } => None, + | ErrorQuickFix::ImportModule { .. } + | 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 c6ccef7079..19b153faa0 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -14,7 +14,6 @@ use pyrefly_python::module::Module; use pyrefly_python::sys_info::PythonVersion; use ruff_python_ast::PySourceType; use ruff_text_size::Ranged; -use pyrefly_python::sys_info::PythonVersion; use ruff_text_size::TextRange; use ruff_text_size::TextSize;