diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index d8dabbf413..a39e61b4ee 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -1,11 +1,15 @@ +use std::borrow::Cow; use std::cmp::Ordering; +use std::collections::BTreeMap; use ruff_db::files::File; use ruff_db::parsed::{ParsedModuleRef, parsed_module}; use ruff_db::source::{SourceText, source_text}; use ruff_diagnostics::Edit; +use ruff_python_ast::StringFlags; use ruff_python_ast::find_node::{CoveringNode, covering_node}; use ruff_python_ast::name::Name; +use ruff_python_ast::str::Quote; use ruff_python_ast::token::{Token, TokenKind, Tokens}; use ruff_python_ast::{self as ast, AnyNodeRef}; use ruff_python_codegen::Stylist; @@ -14,8 +18,8 @@ use rustc_hash::FxHashSet; use ty_module_resolver::{KnownModule, ModuleName}; use ty_python_semantic::types::UnionType; use ty_python_semantic::{ - Completion as SemanticCompletion, NameKind, SemanticModel, - types::{CycleDetector, KnownClass, Type}, + Completion as SemanticCompletion, HasType, NameKind, SemanticModel, + types::{CycleDetector, KnownClass, StringLiteralType, Type}, }; use crate::docstring::Docstring; @@ -46,6 +50,16 @@ pub fn completion<'db>( ContextKind::Import(ref import) => { import.add_completions(db, file, &mut completions); } + ContextKind::StringLiteral(ref string_context) => { + add_string_literal_completions( + db, + file, + &parsed, + &context.cursor, + string_context, + &mut completions, + ); + } ContextKind::NonImport(ref non_import) => { let model = SemanticModel::new(db, file); let (semantic_completions, scoped) = match non_import.target { @@ -377,6 +391,7 @@ impl<'db> Completion<'db> { } }) } + self.kind.or_else(|| { self.ty .and_then(|ty| imp(db, ty, &CompletionKindVisitor::default())) @@ -492,6 +507,7 @@ struct Context<'m> { #[derive(Debug)] enum ContextKind<'m> { Import(ImportStatement<'m>), + StringLiteral(StringLiteralContext<'m>), NonImport(ContextNonImport<'m>), } @@ -502,6 +518,18 @@ struct ContextNonImport<'m> { target: CompletionTargetAst<'m>, } +#[derive(Debug)] +struct StringLiteralContext<'m> { + literal: &'m ast::ExprStringLiteral, + site: StringLiteralSite<'m>, +} + +#[derive(Debug)] +enum StringLiteralSite<'m> { + CallArg, + AnnAssignValue { annotation: &'m ast::Expr }, +} + impl<'m> Context<'m> { /// Create a new context for finding completions. fn new( @@ -516,6 +544,15 @@ impl<'m> Context<'m> { return None; } + if cursor.is_in_string() { + let site = cursor.string_literal_site()?; + let literal = cursor.string_literal()?; + return Some(Context { + kind: ContextKind::StringLiteral(StringLiteralContext { literal, site }), + cursor, + }); + } + let kind = if let Some(import) = ImportStatement::detect(db, file, &cursor) { ContextKind::Import(import) } else { @@ -530,7 +567,7 @@ impl<'m> Context<'m> { /// Returns a filtering context for use with a completion collector. fn collection_context<'db>(&self, db: &'db dyn Db) -> CollectionContext<'db> { match self.kind { - ContextKind::Import(_) => CollectionContext::none(), + ContextKind::Import(_) | ContextKind::StringLiteral(_) => CollectionContext::none(), ContextKind::NonImport(_) => { let is_raising_exception = self.cursor.is_raising_exception(); CollectionContext { @@ -648,7 +685,7 @@ impl<'m> ContextCursor<'m> { /// Whether the last token is in a place where we should not provide completions. fn is_in_no_completions_place(&self) -> bool { - self.is_in_comment() || self.is_in_string() || self.is_in_definition_place() + self.is_in_comment() || self.is_in_definition_place() } /// Whether the last token is within a comment or not. @@ -662,6 +699,54 @@ impl<'m> ContextCursor<'m> { /// /// Note that this will return `false` when the last token is positioned within an /// interpolation block in an f-string or a t-string. + fn string_literal(&self) -> Option<&'m ast::ExprStringLiteral> { + if !self.is_in_string() { + return None; + } + + let range = TextRange::empty(self.offset); + let node = self + .covering_node(range) + .find_last(|node| node.is_expr_string_literal()) + .ok()?; + + let AnyNodeRef::ExprStringLiteral(literal) = node.node() else { + return None; + }; + Some(literal) + } + + /// Determine whether we're in a syntactic place where literal suggestions make sense. + fn string_literal_site(&self) -> Option> { + // Call argument site (and keyword args, I suppose) + if self + .covering_node(TextRange::empty(self.offset)) + .ancestors() + .take_while(|node| !node.is_statement()) + .any(|node| node.is_arguments()) + { + return Some(StringLiteralSite::CallArg); + } + + // Annotated assignment RHS: `x: T = "..."` + let range = TextRange::empty(self.offset); + let covering = self.covering_node(range); + let ann = covering.ancestors().find_map(|node| match node { + AnyNodeRef::StmtAnnAssign(ann_assign) => { + let in_value = ann_assign + .value + .as_deref() + .is_some_and(|value| value.range().contains_range(range)); + in_value.then_some(ann_assign) + } + _ => None, + })?; + + Some(StringLiteralSite::AnnAssignValue { + annotation: &ann.annotation, + }) + } + fn is_in_string(&self) -> bool { self.tokens_before.last().is_some_and(|t| { matches!( @@ -1092,6 +1177,159 @@ enum Sort { Lower, } +fn add_string_literal_completions<'db>( + db: &'db dyn Db, + file: File, + _parsed: &ParsedModuleRef, + cursor: &ContextCursor<'_>, + ctx: &StringLiteralContext<'_>, + completions: &mut Completions<'db>, +) { + let mut out: BTreeMap<&'db str, StringLiteralType<'db>> = BTreeMap::new(); + let mut seen_types: FxHashSet> = FxHashSet::default(); + + match ctx.site { + StringLiteralSite::CallArg => { + collect_call_site_literals(db, file, cursor.offset, &mut out, &mut seen_types); + } + StringLiteralSite::AnnAssignValue { annotation } => { + let model = SemanticModel::new(db, file); + collect_annotation_site_literals(db, &model, annotation, &mut out, &mut seen_types); + } + } + + if out.is_empty() { + return; + } + + let quote_char = match ctx.literal.value.first_literal_flags().quote_style() { + Quote::Single => '\'', + Quote::Double => '"', + }; + + for (value, literal_ty) in out { + let escaped = escape_for_quote(value, quote_char); + let insert: Box = match escaped { + Cow::Borrowed(s) => Box::::from(s), + Cow::Owned(s) => s.into_boxed_str(), + }; + + completions.add(Completion { + name: Name::new(value), + qualified: None, + insert: Some(insert), + ty: Some(Type::StringLiteral(literal_ty)), + kind: Some(CompletionKind::Value), + module_name: None, + import: None, + builtin: false, + is_type_check_only: false, + is_definitively_raisable: false, + documentation: None, + }); + } +} + +fn collect_call_site_literals<'db>( + db: &'db dyn Db, + file: File, + offset: TextSize, + out: &mut BTreeMap<&'db str, StringLiteralType<'db>>, + seen_types: &mut FxHashSet>, +) { + let Some(sig_help) = signature_help(db, file, offset) else { + return; + }; + + if let Some(i) = sig_help.active_signature { + if let Some(sig) = sig_help.signatures.get(i) { + collect_from_signature(db, sig, seen_types, out); + } + return; + } + + for sig in &sig_help.signatures { + collect_from_signature(db, sig, seen_types, out); + } +} + +fn collect_from_signature<'db>( + db: &'db dyn Db, + signature: &crate::SignatureDetails<'db>, + seen_types: &mut FxHashSet>, + out: &mut BTreeMap<&'db str, StringLiteralType<'db>>, +) { + let Some(active_parameter) = signature.active_parameter else { + return; + }; + let Some(parameter) = signature.parameters.get(active_parameter) else { + return; + }; + let Some(ty) = parameter.ty else { + return; + }; + + collect_string_literals_from_type(db, ty, seen_types, out); +} + +fn collect_annotation_site_literals<'db>( + db: &'db dyn Db, + model: &SemanticModel<'db>, + annotation: &ast::Expr, + out: &mut BTreeMap<&'db str, StringLiteralType<'db>>, + seen_types: &mut FxHashSet>, +) { + let Some(annotation_ty) = annotation.inferred_type(model) else { + return; + }; + collect_string_literals_from_type(db, annotation_ty, seen_types, out); +} + +fn collect_string_literals_from_type<'db>( + db: &'db dyn Db, + ty: Type<'db>, + seen_types: &mut FxHashSet>, + out: &mut BTreeMap<&'db str, StringLiteralType<'db>>, +) { + if !seen_types.insert(ty) { + return; + } + + match ty { + Type::StringLiteral(literal) => { + out.entry(literal.as_str(db)).or_insert(literal); + } + Type::Union(union) => { + for ty in union.elements(db) { + collect_string_literals_from_type(db, *ty, seen_types, out); + } + } + Type::Intersection(intersection) => { + for ty in intersection.iter_positive(db) { + collect_string_literals_from_type(db, ty, seen_types, out); + } + } + Type::TypeAlias(alias) => { + collect_string_literals_from_type(db, alias.value_type(db), seen_types, out); + } + _ => {} + } +} + +fn escape_for_quote<'a>(value: &'a str, quote: char) -> Cow<'a, str> { + if !value.bytes().any(|b| b == b'\\' || b == quote as u8) { + return Cow::Borrowed(value); + } + let mut escaped = String::with_capacity(value.len()); + for ch in value.chars() { + if ch == quote || ch == '\\' { + escaped.push('\\'); + } + escaped.push(ch); + } + Cow::Owned(escaped) +} + /// Detect and construct completions for unset function arguments. /// /// Suggestions are only provided if the cursor is currently inside a diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index cc05c9c000..8d96bf255c 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -14502,6 +14502,11 @@ impl<'db> StringLiteralType<'db> { pub(crate) fn python_len(self, db: &'db dyn Db) -> usize { self.value(db).chars().count() } + + /// Returns the contents of this string literal. + pub fn as_str(self, db: &'db dyn Db) -> &'db str { + self.value(db) + } } /// # Ordering diff --git a/crates/ty_server/tests/e2e/completions.rs b/crates/ty_server/tests/e2e/completions.rs index e5a5448027..234c52af9f 100644 --- a/crates/ty_server/tests/e2e/completions.rs +++ b/crates/ty_server/tests/e2e/completions.rs @@ -55,6 +55,153 @@ walktr Ok(()) } +/// Tests that string literal completions are offered for call arguments. +#[test] +fn string_literal_completions_for_calls() -> Result<()> { + let workspace_root = SystemPath::new("src"); + let foo = SystemPath::new("src/foo.py"); + let foo_content = "\ +from typing import Literal + +A = Literal[\"a\", \"b\", \"c\"] +def func(a: A): + ... + +func(\" \") +"; + + let mut server = TestServerBuilder::new()? + .with_initialization_options(ClientOptions::default()) + .with_workspace(workspace_root, None)? + .with_file(foo, foo_content)? + .build() + .wait_until_workspaces_are_initialized(); + + server.open_text_document(foo, foo_content, 1); + let _ = server.await_notification::(); + + let completions = server.completion_request(&server.file_uri(foo), Position::new(6, 6)); + insta::assert_json_snapshot!(completions, @r#" + [ + { + "label": "a", + "kind": 12, + "detail": "Literal[/"a/"]", + "sortText": "0", + "insertText": "a" + }, + { + "label": "b", + "kind": 12, + "detail": "Literal[/"b/"]", + "sortText": "1", + "insertText": "b" + }, + { + "label": "c", + "kind": 12, + "detail": "Literal[/"c/"]", + "sortText": "2", + "insertText": "c" + } + ] + "#); + + Ok(()) +} + +/// Tests that string literal completions are offered when assigning to typed variables. +#[test] +fn string_literal_completions_for_typed_assignment() -> Result<()> { + let workspace_root = SystemPath::new("src"); + let foo = SystemPath::new("src/foo.py"); + let foo_content = "\ +from typing import Literal + +value: Literal[\"x\", \"y\"] = \" \" +"; + + let mut server = TestServerBuilder::new()? + .with_initialization_options(ClientOptions::default()) + .with_workspace(workspace_root, None)? + .with_file(foo, foo_content)? + .build() + .wait_until_workspaces_are_initialized(); + + server.open_text_document(foo, foo_content, 1); + let _ = server.await_notification::(); + + let completions = server.completion_request(&server.file_uri(foo), Position::new(2, 28)); + insta::assert_json_snapshot!(completions, @r#" + [ + { + "label": "x", + "kind": 12, + "detail": "Literal[/"x/"]", + "sortText": "0", + "insertText": "x" + }, + { + "label": "y", + "kind": 12, + "detail": "Literal[/"y/"]", + "sortText": "1", + "insertText": "y" + } + ] + "#); + + Ok(()) +} + +/// Tests that only string literal values are suggested from mixed literal types. +#[test] +fn string_literal_completions_filter_non_strings() -> Result<()> { + let workspace_root = SystemPath::new("src"); + let foo = SystemPath::new("src/foo.py"); + let foo_content = "\ +from typing import Literal + +Mixed = Literal[\"left\", 1, \"right\"] +def consume(value: Mixed): + ... + +consume(\" \") +"; + + let mut server = TestServerBuilder::new()? + .with_initialization_options(ClientOptions::default()) + .with_workspace(workspace_root, None)? + .with_file(foo, foo_content)? + .build() + .wait_until_workspaces_are_initialized(); + + server.open_text_document(foo, foo_content, 1); + let _ = server.await_notification::(); + + let completions = server.completion_request(&server.file_uri(foo), Position::new(6, 9)); + insta::assert_json_snapshot!(completions, @r#" + [ + { + "label": "left", + "kind": 12, + "detail": "Literal[/"left/"]", + "sortText": "0", + "insertText": "left" + }, + { + "label": "right", + "kind": 12, + "detail": "Literal[/"right/"]", + "sortText": "1", + "insertText": "right" + } + ] + "#); + + Ok(()) +} + /// Tests that disabling auto-import works. #[test] fn disable_auto_import() -> Result<()> {