fix autocomplete for discriminated unions #1442#3023
fix autocomplete for discriminated unions #1442#3023asukaminato0721 wants to merge 1 commit intofacebook:mainfrom
Conversation
There was a problem hiding this comment.
Pull request overview
This PR fixes LSP completions for discriminated unions of TypedDict in dict literals by recovering the dict literal’s expected/contextual type, narrowing union members based on already-entered literal fields (e.g. "kind": "foo"), and using the narrowed set to drive both key and literal-value suggestions.
Changes:
- Add TypedDict-union narrowing for dict literals based on already-present literal fields.
- Add dict-literal string value completion for
Literal[...]-typed TypedDict fields (e.g. completing"kind": "|") in dict literals. - Add regression tests covering both discriminated-union value completion and narrowed key completion.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
pyrefly/lib/state/lsp/dict_completions.rs |
Implements expected-type recovery for dict literals, TypedDict-union member narrowing, filters already-present keys, and adds dict value literal completions. |
pyrefly/lib/lsp/wasm/completion.rs |
Wires the new dict value literal completion provider into the completion pipeline. |
pyrefly/lib/test/lsp/completion.rs |
Adds regression tests for discriminated-union TypedDict key/value completion behavior (issue #1442). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| fn typed_dict_field_type_in_member<'b>( | ||
| solver: &crate::alt::answers_solver::AnswersSolver< | ||
| crate::state::lsp::TransactionHandle<'b>, | ||
| >, | ||
| member: &Type, | ||
| key: &str, | ||
| ) -> Option<Type> { | ||
| let typed_dict = match member { | ||
| Type::TypedDict(td) | Type::PartialTypedDict(td) => td, | ||
| _ => return None, | ||
| }; | ||
| solver | ||
| .type_order() | ||
| .typed_dict_fields(typed_dict) | ||
| .iter() | ||
| .find_map(|(name, field)| (name == key).then(|| field.ty.clone())) | ||
| } |
There was a problem hiding this comment.
typed_dict_field_type_in_member linearly scans typed_dict_fields(...) for every (member, item) pair during narrowing. Since completions run on essentially every keystroke, this can become noticeably expensive for large TypedDicts or large unions. Consider building a per-member field-name -> type map once (or otherwise avoiding repeated linear scans) inside the narrowing / key-collection paths.
| fn dict_literal_present_keys( | ||
| dict: &ExprDict, | ||
| skip_key_range: Option<TextRange>, | ||
| ) -> BTreeMap<String, ()> { | ||
| dict.items | ||
| .iter() | ||
| .filter_map(|item| { | ||
| let Expr::StringLiteral(lit) = item.key.as_ref()? else { | ||
| return None; | ||
| }; | ||
| (skip_key_range != Some(lit.range())).then(|| (lit.value.to_string(), ())) | ||
| }) | ||
| .collect() | ||
| } |
There was a problem hiding this comment.
dict_literal_present_keys uses a BTreeMap<String, ()> purely for membership checks. A BTreeSet<String> (or similar set type) would better express intent and avoid storing unused unit values.
|
Hey @asukaminato0721, thanks for putting this together! Discriminated union narrowing for TypedDict is a massive improvement for the LSP experience. I took a look through the logic and I’d like to second the AI's feedback on two points to help get this 'merge-ready': -Performance: In typed_dict_field_type_in_member, the linear scan on typed_dict_fields could definitely become a bottleneck for large projects since completions fire on every keystroke. Mapping these fields once before the narrowing loop would keep the LSP snappy. -Idiomatic Rust: Switching the BTreeMap<String, ()> to a BTreeSet in dict_literal_present_keys would be much cleaner since we only care about membership. Regarding the mypy_primer failure: I noticed check (7) is failing. If you need help debugging whether that’s a regression or just a flake, let me know I’m happy to pull the branch and help investigate! |
|
According to mypy_primer, this change doesn't affect type check results on a corpus of open source code. ✅ |
Summary
Fixes #1442
The completion engine now recovers the enclosing expected type for incomplete dict literals, narrows union members from already-entered literal fields like "kind": "foo", and uses that narrowed set for both key suggestions and literal value suggestions.
Test Plan
add test