From 3eefb22283bbcaa0b7a75aa6aa44c29676711eef Mon Sep 17 00:00:00 2001 From: Iqbal Date: Sat, 1 Aug 2026 18:18:19 +0700 Subject: [PATCH 1/2] [pyrefly] Resolve TypedDict keys in go-to-definition TypedDict string keys currently fall through to operator navigation and resolve to dict.__getitem__. Resolve literal keys against the inferred TypedDict class fields and MRO before that fallback, covering class-based, inherited, and functional declarations while preserving ordinary subscript behavior. --- pyrefly/lib/state/lsp.rs | 63 ++++++++++++++++++++++++++++++ pyrefly/lib/test/lsp/definition.rs | 51 ++++++++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 45139c54e2..63d754cad6 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -8,6 +8,7 @@ use std::cmp::Ordering; use std::cmp::Reverse; use std::collections::HashSet; +use std::iter; use std::path::PathBuf; use std::sync::Arc; use std::sync::LazyLock; @@ -32,6 +33,7 @@ use pyrefly_python::short_identifier::ShortIdentifier; use pyrefly_python::symbol_kind::SymbolKind; use pyrefly_python::sys_info::SysInfo; use pyrefly_types::type_alias::TypeAliasData; +use pyrefly_types::typed_dict::TypedDict; use pyrefly_util::gas::Gas; use pyrefly_util::lock::Mutex; use pyrefly_util::prelude::SliceExt; @@ -2303,6 +2305,62 @@ impl<'a> Transaction<'a> { Ok(Some(defs)) } + /// Resolve a string subscript on a TypedDict to the field declaration. + fn find_definition_for_typed_dict_key( + &self, + handle: &Handle, + position: TextSize, + covering_nodes: &[AnyNodeRef], + ) -> Option { + let subscript = covering_nodes.iter().find_map(|node| match node { + AnyNodeRef::ExprSubscript(subscript) => Some(subscript), + _ => None, + })?; + let Expr::StringLiteral(key) = subscript.slice.as_ref() else { + return None; + }; + if !key.range().contains(position) { + return None; + } + + let name = Name::new(key.value.to_str()); + let base_type = self.get_type_trace(handle, subscript.value.range())?; + let typed_dict = match base_type { + Type::TypedDict(TypedDict::TypedDict(typed_dict)) + | Type::PartialTypedDict(TypedDict::TypedDict(typed_dict)) => typed_dict, + _ => return None, + }; + let display_name = name.to_string(); + let (module, definition_range, docstring_range) = self + .ad_hoc_solve(handle, "typed_dict_key_definition", |solver| { + let class = typed_dict.class_object(); + let mro = solver.get_mro_for_class(class); + iter::once(class) + .chain( + mro.ancestors_no_object() + .iter() + .map(|ancestor| ancestor.class_object()), + ) + .find_map(|class| { + let fields = solver.get_class_fields(class)?; + Some(( + class.module().dupe(), + fields.field_decl_range(&name)?, + fields.field_docstring_range(&name), + )) + }) + }) + .flatten()?; + + Some(FindDefinitionItemWithDocstring { + metadata: DefinitionMetadata::Attribute, + definition_range, + module, + docstring_range, + display_name: Some(display_name), + }) + } + pub fn find_definition_for_attribute( &self, handle: &Handle, @@ -2841,6 +2899,11 @@ impl<'a> Transaction<'a> { }), }; } + if let Some(definition) = + self.find_definition_for_typed_dict_key(handle, position, &covering_nodes) + { + return Ok(vec1![definition]); + } // Fall back to operator handling if let Some(defs) = self.find_definition_for_operator(handle, &covering_nodes, preference)? diff --git a/pyrefly/lib/test/lsp/definition.rs b/pyrefly/lib/test/lsp/definition.rs index 555030e69c..35f61817c8 100644 --- a/pyrefly/lib/test/lsp/definition.rs +++ b/pyrefly/lib/test/lsp/definition.rs @@ -162,6 +162,57 @@ Definition Result: ); } +#[test] +fn typed_dict_key_test() { + let code = r#" +from typing import TypedDict + +class Person(TypedDict): + name: str + +class Employee(Person): + employee_id: int + +Functional = TypedDict("Functional", {"title": str}) + +person: Person = {"name": ""} +employee: Employee = {"name": "", "employee_id": 0} +functional: Functional = {"title": ""} + +person["name"] +# ^ +employee["name"] +# ^ +functional["title"] +# ^ +"#; + let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report); + assert_eq!( + r#" +# main.py +16 | person["name"] + ^ +Definition Result: +5 | name: str + ^^^^ + +18 | employee["name"] + ^ +Definition Result: +5 | name: str + ^^^^ + +20 | functional["title"] + ^ +Definition Result: +10 | Functional = TypedDict("Functional", {"title": str}) + ^^^^^^^ +"# + .trim(), + report.trim(), + ); +} + #[test] fn pytest_fixture_parameter_goes_to_fixture_definition() { let code = r#" From 59a69cdfcc22c262ad4ac82d7b137d1f27e3809d Mon Sep 17 00:00:00 2001 From: Iqbal Date: Sun, 2 Aug 2026 13:43:24 +0700 Subject: [PATCH 2/2] [pyrefly] Stop definition fallback for missing TypedDict keys Distinguish declared TypedDict key access from ordinary subscripts before resolving the field. This lets missing fields return no definition instead of falling through to __getitem__, while keeping the field lookup itself optional and covering both class and functional declarations. --- pyrefly/lib/state/lsp.rs | 83 ++++++++++++++++++------------ pyrefly/lib/test/lsp/definition.rs | 35 +++++++++++++ 2 files changed, 85 insertions(+), 33 deletions(-) diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 63d754cad6..ebaea953c2 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -34,6 +34,7 @@ use pyrefly_python::symbol_kind::SymbolKind; use pyrefly_python::sys_info::SysInfo; use pyrefly_types::type_alias::TypeAliasData; use pyrefly_types::typed_dict::TypedDict; +use pyrefly_types::typed_dict::TypedDictInner; use pyrefly_util::gas::Gas; use pyrefly_util::lock::Mutex; use pyrefly_util::prelude::SliceExt; @@ -2305,13 +2306,13 @@ impl<'a> Transaction<'a> { Ok(Some(defs)) } - /// Resolve a string subscript on a TypedDict to the field declaration. - fn find_definition_for_typed_dict_key( + /// Return the declared TypedDict and key when the cursor is on a string subscript. + fn typed_dict_key_at( &self, handle: &Handle, position: TextSize, covering_nodes: &[AnyNodeRef], - ) -> Option { + ) -> Option<(TypedDictInner, Name)> { let subscript = covering_nodes.iter().find_map(|node| match node { AnyNodeRef::ExprSubscript(subscript) => Some(subscript), _ => None, @@ -2330,35 +2331,44 @@ impl<'a> Transaction<'a> { | Type::PartialTypedDict(TypedDict::TypedDict(typed_dict)) => typed_dict, _ => return None, }; - let display_name = name.to_string(); - let (module, definition_range, docstring_range) = self - .ad_hoc_solve(handle, "typed_dict_key_definition", |solver| { - let class = typed_dict.class_object(); - let mro = solver.get_mro_for_class(class); - iter::once(class) - .chain( - mro.ancestors_no_object() - .iter() - .map(|ancestor| ancestor.class_object()), - ) - .find_map(|class| { - let fields = solver.get_class_fields(class)?; - Some(( - class.module().dupe(), - fields.field_decl_range(&name)?, - fields.field_docstring_range(&name), - )) - }) - }) - .flatten()?; + Some((typed_dict, name)) + } - Some(FindDefinitionItemWithDocstring { - metadata: DefinitionMetadata::Attribute, - definition_range, - module, - docstring_range, - display_name: Some(display_name), + /// Resolve a TypedDict key to its field declaration. + fn find_definition_for_typed_dict_key( + &self, + handle: &Handle, + typed_dict: &TypedDictInner, + name: &Name, + ) -> Option { + self.ad_hoc_solve(handle, "typed_dict_key_definition", |solver| { + let class = typed_dict.class_object(); + let mro = solver.get_mro_for_class(class); + iter::once(class) + .chain( + mro.ancestors_no_object() + .iter() + .map(|ancestor| ancestor.class_object()), + ) + .find_map(|class| { + let fields = solver.get_class_fields(class)?; + Some(( + class.module().dupe(), + fields.field_decl_range(name)?, + fields.field_docstring_range(name), + )) + }) }) + .flatten() + .map( + |(module, definition_range, docstring_range)| FindDefinitionItemWithDocstring { + metadata: DefinitionMetadata::Attribute, + definition_range, + module, + docstring_range, + display_name: Some(name.to_string()), + }, + ) } pub fn find_definition_for_attribute( @@ -2899,10 +2909,17 @@ impl<'a> Transaction<'a> { }), }; } - if let Some(definition) = - self.find_definition_for_typed_dict_key(handle, position, &covering_nodes) + if let Some((typed_dict, name)) = + self.typed_dict_key_at(handle, position, &covering_nodes) { - return Ok(vec1![definition]); + return match self.find_definition_for_typed_dict_key(handle, &typed_dict, &name) + { + Some(definition) => Ok(vec1![definition]), + None => Err(EmptyResponseReason::DefinitionNotFound { + name: name.to_string(), + context: DefinitionContext::Attribute, + }), + }; } // Fall back to operator handling if let Some(defs) = diff --git a/pyrefly/lib/test/lsp/definition.rs b/pyrefly/lib/test/lsp/definition.rs index 35f61817c8..be4c79d44a 100644 --- a/pyrefly/lib/test/lsp/definition.rs +++ b/pyrefly/lib/test/lsp/definition.rs @@ -213,6 +213,41 @@ Definition Result: ); } +#[test] +fn typed_dict_missing_key_test() { + let code = r#" +from typing import TypedDict + +class Person(TypedDict): + name: str + +Functional = TypedDict("Functional", {"title": str}) + +person: Person = {"name": ""} +functional: Functional = {"title": ""} + +person["missing"] +# ^ +functional["missing"] +# ^ +"#; + let report = get_batched_lsp_operations_report_allow_error(&[("main", code)], get_test_report); + assert_eq!( + r#" +# main.py +12 | person["missing"] + ^ +Definition Result: None + +14 | functional["missing"] + ^ +Definition Result: None +"# + .trim(), + report.trim(), + ); +} + #[test] fn pytest_fixture_parameter_goes_to_fixture_definition() { let code = r#"