[ty] Track dictionary literal keys as individual places

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
+```

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:

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<Item = &'ast ast::Expr> + 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,18 +2536,26 @@ 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,
                                         value: &node.value,
                                         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);

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<AnnotatedAssignmentDefinitionNodeRef<'ast>> for DefinitionNodeRe
     }
 }
 
+impl<'ast, 'db> From<DictKeyAssignmentNodeRef<'ast, 'db>> for DefinitionNodeRef<'ast, 'db> {
+    fn from(node_ref: DictKeyAssignmentNodeRef<'ast, 'db>) -> Self {
+        Self::DictKeyAssignment(node_ref)
+    }
+}
+
 impl<'ast, 'db> From<WithItemDefinitionNodeRef<'ast, 'db>> 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<ast::StmtAugAssign>),
+    DictKeyAssignment(DictKeyAssignmentKind<'db>),
     For(ForStmtDefinitionKind<'db>),
     Comprehension(ComprehensionDefinitionKind<'db>),
     VariadicPositionalParameter(AstNodeRef<ast::Parameter>),
@@ -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<ast::Expr>,
+    pub(crate) value: AstNodeRef<ast::Expr>,
+    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>,