diff --git a/crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md b/crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md index 8c0921ccd0ec20..5716d0415053c7 100644 --- a/crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md +++ b/crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md @@ -75,3 +75,45 @@ reveal_type({"a": 1, "b": (1, 2), "c": (1, 2, 3)}) # revealed: dict[Unknown | int, Unknown | int] reveal_type({x: y for x, y in enumerate(range(42))}) ``` + +## Key narrowing + +The original assignment to each key, as well as future assignments, are used to narrow access to +individual keys: + +```py +from typing import TypedDict + +x1 = {"a": 1, "b": "2"} +reveal_type(x1) # revealed: dict[Unknown | str, Unknown | int | str] +reveal_type(x1["a"]) # revealed: Literal[1] +reveal_type(x1["b"]) # revealed: Literal["2"] + +x1["a"] = 2 +reveal_type(x1["a"]) # revealed: Literal[2] + +x2: dict[str, int | str] = {"a": 1, "b": "2"} +reveal_type(x2) # revealed: dict[str, int | str] +reveal_type(x2["a"]) # revealed: Literal[1] +reveal_type(x2["b"]) # revealed: Literal["2"] + +class TD(TypedDict): + td: int + +x3: dict[int, int | TD] = {1: 1, 2: {"td": 1}} +reveal_type(x3) # revealed: dict[int, int | TD] +reveal_type(x3[1]) # revealed: Literal[1] +reveal_type(x3[2]) # revealed: TD + +x4 = {"a": 1, "b": {"c": 2, "d": "3"}} +reveal_type(x4["a"]) # revealed: Literal[1] +reveal_type(x4["b"]) # revealed: dict[Unknown | str, Unknown | int | str] +reveal_type(x4["b"]["c"]) # revealed: Literal[2] +reveal_type(x4["b"]["d"]) # revealed: Literal["3"] + +x5: dict[str, int | dict[str, int | TD]] = {"a": 1, "b": {"c": 2, "d": {"td": 1}}} +reveal_type(x5["a"]) # revealed: Literal[1] +reveal_type(x5["b"]) # revealed: dict[str, int | TD] +reveal_type(x5["b"]["c"]) # revealed: Literal[2] +reveal_type(x5["b"]["d"]) # revealed: TD +``` diff --git a/crates/ty_python_semantic/resources/mdtest/typed_dict.md b/crates/ty_python_semantic/resources/mdtest/typed_dict.md index 36e4a3a58e0f04..019d05a8ce1589 100644 --- a/crates/ty_python_semantic/resources/mdtest/typed_dict.md +++ b/crates/ty_python_semantic/resources/mdtest/typed_dict.md @@ -26,8 +26,8 @@ inferred based on the `TypedDict` definition: ```py alice: Person = {"name": "Alice", "age": 30} -reveal_type(alice["name"]) # revealed: str -reveal_type(alice["age"]) # revealed: int | None +reveal_type(alice["name"]) # revealed: Literal["Alice"] +reveal_type(alice["age"]) # revealed: Literal[30] # error: [invalid-key] "Unknown key "non_existing" for TypedDict `Person`" reveal_type(alice["non_existing"]) # revealed: Unknown @@ -140,7 +140,7 @@ reveal_type(plot2["y"]) # revealed: list[int | None] plot3: Plot = {"y": homogeneous_list(1, 2, 3), "x": homogeneous_list(1, 2, 3)} reveal_type(plot3["y"]) # revealed: list[int | None] -reveal_type(plot3["x"]) # revealed: list[int | None] | None +reveal_type(plot3["x"]) # revealed: list[int | None] Y = "y" X = "x" @@ -194,8 +194,8 @@ class Person(TypedDict): ```py alice: Person = {"inner": {"name": "Alice", "age": 30}} -reveal_type(alice["inner"]["name"]) # revealed: str -reveal_type(alice["inner"]["age"]) # revealed: int | None +reveal_type(alice["inner"]["name"]) # revealed: Literal["Alice"] +reveal_type(alice["inner"]["age"]) # revealed: Literal[30] # error: [invalid-key] "Unknown key "non_existing" for TypedDict `Inner`" reveal_type(alice["inner"]["non_existing"]) # revealed: Unknown @@ -778,7 +778,7 @@ alice: Person = {"name": "Alice"} # error: [invalid-argument-type] "Argument to function `dangerous` is incorrect: Expected `dict[str, object]`, found `Person`" dangerous(alice) -reveal_type(alice["name"]) # revealed: str +reveal_type(alice["name"]) # revealed: Literal["Alice"] ``` Likewise, `dict`s are not assignable to typed dictionaries: diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index 82a38b10de8d2c..7dabf316b44738 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -25,11 +25,13 @@ use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey; use crate::semantic_index::definition::{ AnnotatedAssignmentDefinitionNodeRef, AssignmentDefinitionNodeRef, ComprehensionDefinitionNodeRef, Definition, DefinitionCategory, DefinitionNodeKey, - DefinitionNodeRef, Definitions, ExceptHandlerDefinitionNodeRef, ForStmtDefinitionNodeRef, - ImportDefinitionNodeRef, ImportFromDefinitionNodeRef, ImportFromSubmoduleDefinitionNodeRef, - MatchPatternDefinitionNodeRef, StarImportDefinitionNodeRef, WithItemDefinitionNodeRef, + DefinitionNodeRef, Definitions, DictKeyAssignmentNodeRef, ExceptHandlerDefinitionNodeRef, + ForStmtDefinitionNodeRef, ImportDefinitionNodeRef, ImportFromDefinitionNodeRef, + ImportFromSubmoduleDefinitionNodeRef, MatchPatternDefinitionNodeRef, + StarImportDefinitionNodeRef, WithItemDefinitionNodeRef, }; use crate::semantic_index::expression::{Expression, ExpressionKind}; +use crate::semantic_index::member::MemberExprBuilder; use crate::semantic_index::place::{PlaceExpr, PlaceTableBuilder, ScopedPlaceId}; use crate::semantic_index::predicate::{ CallableAndCallExpr, ClassPatternKind, PatternPredicate, PatternPredicateKind, Predicate, @@ -756,6 +758,59 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { (definition, num_definitions) } + // Creates a definition for each key-value assignment in the dictionary. + // + // If there are multiple targets, a given key-value definition will be created multiple + // times for each target. + fn add_dict_key_assignment_definitions( + &mut self, + targets: impl IntoIterator + Copy, + dict: &'ast ast::ExprDict, + assignment: Definition<'db>, + ) { + for target in targets { + if let Some(target) = MemberExprBuilder::visit_expr(target.into()) { + self.add_dict_key_assignment_definitions_impl(&target, dict, assignment); + } + } + } + + fn add_dict_key_assignment_definitions_impl( + &mut self, + target: &MemberExprBuilder, + dict: &'ast ast::ExprDict, + assignment: Definition<'db>, + ) { + for item in &dict.items { + let Some(key) = item.key.as_ref() else { + continue; + }; + + let Some(member_expr) = MemberExprBuilder::visit_subscript_expr(target.clone(), key) + else { + continue; + }; + + // Recurse into nested dictionaries. + if let ast::Expr::Dict(dict_value) = &item.value { + self.add_dict_key_assignment_definitions_impl(&member_expr, dict_value, assignment); + } + + if let Some(place_expr) = PlaceExpr::try_from_member_expr(member_expr) { + let place_id = self.add_place(place_expr); + + self.add_definition( + place_id, + DictKeyAssignmentNodeRef { + key, + assignment, + value: &item.value, + }, + ); + } + } + } + fn record_expression_narrowing_constraint( &mut self, predicate_node: &ast::Expr, @@ -2481,7 +2536,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { if is_definition { match self.current_assignment() { Some(CurrentAssignment::Assign { node, unpack }) => { - self.add_definition( + let assignment = self.add_definition( place_id, AssignmentDefinitionNodeRef { unpack, @@ -2489,10 +2544,18 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { target: expr, }, ); + + if let ast::Expr::Dict(dict) = &*node.value { + self.add_dict_key_assignment_definitions( + &node.targets, + dict, + assignment, + ); + } } Some(CurrentAssignment::AnnAssign(ann_assign)) => { self.add_standalone_type_expression(&ann_assign.annotation); - self.add_definition( + let assignment = self.add_definition( place_id, AnnotatedAssignmentDefinitionNodeRef { node: ann_assign, @@ -2501,6 +2564,14 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { target: expr, }, ); + + if let Some(ast::Expr::Dict(dict)) = ann_assign.value.as_deref() { + self.add_dict_key_assignment_definitions( + [&*ann_assign.target], + dict, + assignment, + ); + } } Some(CurrentAssignment::AugAssign(aug_assign)) => { self.add_definition(place_id, aug_assign); diff --git a/crates/ty_python_semantic/src/semantic_index/definition.rs b/crates/ty_python_semantic/src/semantic_index/definition.rs index e0c7cf8274c0a5..218e819cfda41a 100644 --- a/crates/ty_python_semantic/src/semantic_index/definition.rs +++ b/crates/ty_python_semantic/src/semantic_index/definition.rs @@ -275,6 +275,7 @@ pub(crate) enum DefinitionNodeRef<'ast, 'db> { Assignment(AssignmentDefinitionNodeRef<'ast, 'db>), AnnotatedAssignment(AnnotatedAssignmentDefinitionNodeRef<'ast>), AugmentedAssignment(&'ast ast::StmtAugAssign), + DictKeyAssignment(DictKeyAssignmentNodeRef<'ast, 'db>), Comprehension(ComprehensionDefinitionNodeRef<'ast, 'db>), VariadicPositionalParameter(&'ast ast::Parameter), VariadicKeywordParameter(&'ast ast::Parameter), @@ -371,6 +372,12 @@ impl<'ast> From> for DefinitionNodeRe } } +impl<'ast, 'db> From> for DefinitionNodeRef<'ast, 'db> { + fn from(node_ref: DictKeyAssignmentNodeRef<'ast, 'db>) -> Self { + Self::DictKeyAssignment(node_ref) + } +} + impl<'ast, 'db> From> for DefinitionNodeRef<'ast, 'db> { fn from(node_ref: WithItemDefinitionNodeRef<'ast, 'db>) -> Self { Self::WithItem(node_ref) @@ -441,6 +448,13 @@ pub(crate) struct AnnotatedAssignmentDefinitionNodeRef<'ast> { pub(crate) target: &'ast ast::Expr, } +#[derive(Copy, Clone, Debug)] +pub(crate) struct DictKeyAssignmentNodeRef<'ast, 'db> { + pub(crate) key: &'ast ast::Expr, + pub(crate) value: &'ast ast::Expr, + pub(crate) assignment: Definition<'db>, +} + #[derive(Copy, Clone, Debug)] pub(crate) struct WithItemDefinitionNodeRef<'ast, 'db> { pub(crate) unpack: Option<(UnpackPosition, Unpack<'db>)>, @@ -550,6 +564,15 @@ impl<'db> DefinitionNodeRef<'_, 'db> { DefinitionNodeRef::AugmentedAssignment(augmented_assignment) => { DefinitionKind::AugmentedAssignment(AstNodeRef::new(parsed, augmented_assignment)) } + DefinitionNodeRef::DictKeyAssignment(DictKeyAssignmentNodeRef { + key, + value, + assignment, + }) => DefinitionKind::DictKeyAssignment(DictKeyAssignmentKind { + key: AstNodeRef::new(parsed, key), + value: AstNodeRef::new(parsed, value), + assignment, + }), DefinitionNodeRef::For(ForStmtDefinitionNodeRef { unpack, iterable, @@ -658,6 +681,7 @@ impl<'db> DefinitionNodeRef<'_, 'db> { }) => DefinitionNodeKey(NodeKey::from_node(target)), Self::AnnotatedAssignment(ann_assign) => ann_assign.node.into(), Self::AugmentedAssignment(node) => node.into(), + Self::DictKeyAssignment(node) => DefinitionNodeKey(NodeKey::from_node(node.key)), Self::For(ForStmtDefinitionNodeRef { target, iterable: _, @@ -742,6 +766,7 @@ pub enum DefinitionKind<'db> { Assignment(AssignmentDefinitionKind<'db>), AnnotatedAssignment(AnnotatedAssignmentDefinitionKind), AugmentedAssignment(AstNodeRef), + DictKeyAssignment(DictKeyAssignmentKind<'db>), For(ForStmtDefinitionKind<'db>), Comprehension(ComprehensionDefinitionKind<'db>), VariadicPositionalParameter(AstNodeRef), @@ -816,6 +841,9 @@ impl DefinitionKind<'_> { DefinitionKind::AugmentedAssignment(aug_assign) => { aug_assign.node(module).target.range() } + DefinitionKind::DictKeyAssignment(dict_key_assignment) => { + dict_key_assignment.key.node(module).range() + } DefinitionKind::For(for_stmt) => for_stmt.target.node(module).range(), DefinitionKind::Comprehension(comp) => comp.target(module).range(), DefinitionKind::VariadicPositionalParameter(parameter) => { @@ -865,6 +893,9 @@ impl DefinitionKind<'_> { full_range } DefinitionKind::AugmentedAssignment(aug_assign) => aug_assign.node(module).range(), + DefinitionKind::DictKeyAssignment(dict_key_assignment) => { + dict_key_assignment.key.node(module).range() + } DefinitionKind::For(for_stmt) => for_stmt.target.node(module).range(), DefinitionKind::Comprehension(comp) => comp.target(module).range(), DefinitionKind::VariadicPositionalParameter(parameter) => { @@ -927,7 +958,8 @@ impl DefinitionKind<'_> { } } // all of these bind values without declaring a type - DefinitionKind::NamedExpression(_) + DefinitionKind::DictKeyAssignment(_) + | DefinitionKind::NamedExpression(_) | DefinitionKind::Assignment(_) | DefinitionKind::AugmentedAssignment(_) | DefinitionKind::For(_) @@ -1145,6 +1177,23 @@ impl AnnotatedAssignmentDefinitionKind { } } +#[derive(Clone, Debug, get_size2::GetSize)] +pub struct DictKeyAssignmentKind<'db> { + pub(crate) key: AstNodeRef, + pub(crate) value: AstNodeRef, + pub(crate) assignment: Definition<'db>, +} + +impl DictKeyAssignmentKind<'_> { + pub(crate) fn key<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::Expr { + self.key.node(module) + } + + pub(crate) fn value<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::Expr { + self.value.node(module) + } +} + #[derive(Clone, Debug, get_size2::GetSize)] pub struct WithItemDefinitionKind<'db> { target_kind: TargetKind<'db>, diff --git a/crates/ty_python_semantic/src/semantic_index/member.rs b/crates/ty_python_semantic/src/semantic_index/member.rs index 1e25c679ce01df..05c4f13d676421 100644 --- a/crates/ty_python_semantic/src/semantic_index/member.rs +++ b/crates/ty_python_semantic/src/semantic_index/member.rs @@ -1,10 +1,13 @@ -use bitflags::bitflags; -use hashbrown::hash_table::Entry; use ruff_index::{IndexVec, newtype_index}; use ruff_python_ast::{self as ast, name::Name}; use ruff_text_size::{TextLen as _, TextRange, TextSize}; + +use bitflags::bitflags; +use hashbrown::hash_table::Entry; use rustc_hash::FxHasher; use smallvec::SmallVec; + +use std::fmt::Write as _; use std::hash::{Hash, Hasher as _}; use std::ops::{Deref, DerefMut}; @@ -162,111 +165,18 @@ pub(crate) struct MemberExpr { } impl MemberExpr { + #[cfg(test)] pub(super) fn try_from_expr(expression: ast::ExprRef<'_>) -> Option { - fn visit(expr: ast::ExprRef) -> Option<(Name, SmallVec<[SegmentInfo; 8]>)> { - use std::fmt::Write as _; - - match expr { - ast::ExprRef::Name(name) => { - Some((name.id.clone(), smallvec::SmallVec::new_const())) - } - ast::ExprRef::Attribute(attribute) => { - let (mut path, mut segments) = visit(ast::ExprRef::from(&attribute.value))?; - - let start_offset = path.text_len(); - let _ = write!(path, "{}", attribute.attr.id); - segments.push(SegmentInfo::new(SegmentKind::Attribute, start_offset)); - - Some((path, segments)) - } - ast::ExprRef::Subscript(subscript) => { - let (mut path, mut segments) = visit((&subscript.value).into())?; - let start_offset = path.text_len(); - - match &*subscript.slice { - // Handle integer subscripts, like `x[0]`. - ast::Expr::NumberLiteral(ast::ExprNumberLiteral { - value: ast::Number::Int(index), - .. - }) => { - let _ = write!(path, "{index}"); - segments - .push(SegmentInfo::new(SegmentKind::IntSubscript, start_offset)); - } - // Handle negative integer subscripts, like `x[-1]`. - ast::Expr::UnaryOp(ast::ExprUnaryOp { - op: ast::UnaryOp::USub, - operand, - .. - }) => match operand.as_ref() { - ast::Expr::NumberLiteral(ast::ExprNumberLiteral { - value: ast::Number::Int(index), - .. - }) => { - let _ = write!(path, "-{index}"); - segments.push(SegmentInfo::new( - SegmentKind::IntSubscript, - start_offset, - )); - } - _ => return None, - }, - // Handle positive integer subscripts with explicit plus, like `x[+1]`. - ast::Expr::UnaryOp(ast::ExprUnaryOp { - op: ast::UnaryOp::UAdd, - operand, - .. - }) => match operand.as_ref() { - ast::Expr::NumberLiteral(ast::ExprNumberLiteral { - value: ast::Number::Int(index), - .. - }) => { - let _ = write!(path, "{index}"); - segments.push(SegmentInfo::new( - SegmentKind::IntSubscript, - start_offset, - )); - } - _ => return None, - }, - // Handle boolean subscripts, like `x[True]` or `x[False]`. - // In Python, `True` and `False` are equivalent to `1` and `0` for indexing. - ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => { - let _ = write!(path, "{}", u8::from(*value)); - segments - .push(SegmentInfo::new(SegmentKind::IntSubscript, start_offset)); - } - ast::Expr::StringLiteral(string) => { - let _ = write!(path, "{}", string.value); - segments - .push(SegmentInfo::new(SegmentKind::StringSubscript, start_offset)); - } - // Handle bytes literal subscripts, like `x[b"key"]`. - ast::Expr::BytesLiteral(bytes) => { - let bytes_vec: Vec = bytes.value.bytes().collect(); - let _ = write!(path, "{}", String::from_utf8_lossy(&bytes_vec)); - segments - .push(SegmentInfo::new(SegmentKind::BytesSubscript, start_offset)); - } - _ => { - return None; - } - } - - Some((path, segments)) - } - _ => None, - } - } - - let (path, segments) = visit(expression)?; + MemberExprBuilder::visit_expr(expression).and_then(Self::try_from_builder) + } - if segments.is_empty() { + pub(super) fn try_from_builder(builder: MemberExprBuilder) -> Option { + if builder.segments.is_empty() { None } else { Some(Self { - path, - segments: Segments::from_vec(segments), + path: builder.path, + segments: Segments::from_vec(builder.segments), }) } } @@ -302,6 +212,116 @@ impl MemberExpr { } } +/// A builder for a [`MemberExpr`]. +#[derive(Clone, Debug, PartialEq, Eq, get_size2::GetSize)] +pub(super) struct MemberExprBuilder { + path: Name, + segments: SmallVec<[SegmentInfo; 8]>, +} + +impl MemberExprBuilder { + pub(super) fn visit_expr(expr: ast::ExprRef) -> Option { + match expr { + ast::ExprRef::Name(name) => Some(MemberExprBuilder { + path: name.id.clone(), + segments: smallvec::SmallVec::new_const(), + }), + + ast::ExprRef::Attribute(attribute) => { + let mut builder = + MemberExprBuilder::visit_expr(ast::ExprRef::from(&attribute.value))?; + + let start_offset = builder.path.text_len(); + let _ = write!(builder.path, "{}", attribute.attr.id); + builder + .segments + .push(SegmentInfo::new(SegmentKind::Attribute, start_offset)); + + Some(builder) + } + ast::ExprRef::Subscript(subscript) => { + let subscript_value = + MemberExprBuilder::visit_expr(ast::ExprRef::from(&subscript.value))?; + MemberExprBuilder::visit_subscript_expr(subscript_value, &subscript.slice) + } + _ => None, + } + } + + pub(super) fn visit_subscript_expr( + subscript_value: MemberExprBuilder, + subscript_slice: &ast::Expr, + ) -> Option { + let MemberExprBuilder { + mut path, + mut segments, + } = subscript_value; + let start_offset = path.text_len(); + + match subscript_slice { + // Handle integer subscripts, like `x[0]`. + ast::Expr::NumberLiteral(ast::ExprNumberLiteral { + value: ast::Number::Int(index), + .. + }) => { + let _ = write!(path, "{index}"); + segments.push(SegmentInfo::new(SegmentKind::IntSubscript, start_offset)); + } + // Handle negative integer subscripts, like `x[-1]`. + ast::Expr::UnaryOp(ast::ExprUnaryOp { + op: ast::UnaryOp::USub, + operand, + .. + }) => match operand.as_ref() { + ast::Expr::NumberLiteral(ast::ExprNumberLiteral { + value: ast::Number::Int(index), + .. + }) => { + let _ = write!(path, "-{index}"); + segments.push(SegmentInfo::new(SegmentKind::IntSubscript, start_offset)); + } + _ => return None, + }, + // Handle positive integer subscripts with explicit plus, like `x[+1]`. + ast::Expr::UnaryOp(ast::ExprUnaryOp { + op: ast::UnaryOp::UAdd, + operand, + .. + }) => match operand.as_ref() { + ast::Expr::NumberLiteral(ast::ExprNumberLiteral { + value: ast::Number::Int(index), + .. + }) => { + let _ = write!(path, "{index}"); + segments.push(SegmentInfo::new(SegmentKind::IntSubscript, start_offset)); + } + _ => return None, + }, + // Handle boolean subscripts, like `x[True]` or `x[False]`. + // In Python, `True` and `False` are equivalent to `1` and `0` for indexing. + ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => { + let _ = write!(path, "{}", u8::from(*value)); + segments.push(SegmentInfo::new(SegmentKind::IntSubscript, start_offset)); + } + ast::Expr::StringLiteral(string) => { + let _ = write!(path, "{}", string.value); + segments.push(SegmentInfo::new(SegmentKind::StringSubscript, start_offset)); + } + // Handle bytes literal subscripts, like `x[b"key"]`. + ast::Expr::BytesLiteral(bytes) => { + let bytes_vec: Vec = bytes.value.bytes().collect(); + let _ = write!(path, "{}", String::from_utf8_lossy(&bytes_vec)); + segments.push(SegmentInfo::new(SegmentKind::BytesSubscript, start_offset)); + } + _ => { + return None; + } + } + + Some(MemberExprBuilder { path, segments }) + } +} + impl std::fmt::Display for MemberExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str(self.symbol_name())?; diff --git a/crates/ty_python_semantic/src/semantic_index/place.rs b/crates/ty_python_semantic/src/semantic_index/place.rs index 850a43bf80a085..1106c5cb11025d 100644 --- a/crates/ty_python_semantic/src/semantic_index/place.rs +++ b/crates/ty_python_semantic/src/semantic_index/place.rs @@ -1,5 +1,6 @@ use crate::semantic_index::member::{ - Member, MemberExpr, MemberExprRef, MemberTable, MemberTableBuilder, ScopedMemberId, + Member, MemberExpr, MemberExprBuilder, MemberExprRef, MemberTable, MemberTableBuilder, + ScopedMemberId, }; use crate::semantic_index::scope::FileScopeId; use crate::semantic_index::symbol::{ScopedSymbolId, Symbol, SymbolTable, SymbolTableBuilder}; @@ -48,7 +49,14 @@ impl PlaceExpr { return Some(PlaceExpr::Symbol(Symbol::new(name.id.clone()))); } - let member_expression = MemberExpr::try_from_expr(expr)?; + MemberExprBuilder::visit_expr(expr).and_then(Self::try_from_member_expr) + } + + /// Tries to create a `PlaceExpr` from a member expression. + /// + /// Returns `None` if the expression is not a valid place expression and `Some` otherwise. + pub(super) fn try_from_member_expr(builder: MemberExprBuilder) -> Option { + let member_expression = MemberExpr::try_from_builder(builder)?; Some(Self::Member(Member::new(member_expression))) } } diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index f52d0108baa123..7b4fdb515a36bd 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/crates/ty_python_semantic/src/types/ide_support.rs @@ -1409,6 +1409,7 @@ mod resolve_definition { | DefinitionKind::Assignment(_) | DefinitionKind::AnnotatedAssignment(_) | DefinitionKind::AugmentedAssignment(_) + | DefinitionKind::DictKeyAssignment(_) | DefinitionKind::For(_) | DefinitionKind::Comprehension(_) | DefinitionKind::VariadicPositionalParameter(_) diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index d65c30d1e9583e..ef596dd9749645 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -1811,6 +1811,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { definition, ); } + DefinitionKind::DictKeyAssignment(dict_key_assignment) => { + self.infer_dict_key_assignment_definition( + dict_key_assignment.key(self.module()), + dict_key_assignment.value(self.module()), + dict_key_assignment.assignment, + definition, + ); + } DefinitionKind::For(for_statement_definition) => { self.infer_for_statement_definition(for_statement_definition, definition); } @@ -8174,6 +8182,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { self.infer_augmented_op(assignment, target_type, value_type) } + fn infer_dict_key_assignment_definition( + &mut self, + key: &'ast ast::Expr, + value: &'ast ast::Expr, + assignment: Definition<'db>, + definition: Definition<'db>, + ) { + let value_ty = infer_definition_types(self.db(), assignment).expression_type(value); + self.add_binding(key.into(), definition) + .insert(self, value_ty); + } + fn infer_type_alias_statement(&mut self, node: &ast::StmtTypeAlias) { self.infer_definition(node); }