[ty] Infer types for key-based access on TypedDicts

crates/ty_python_semantic/Cargo.toml

@@ -48,6 +48,7 @@ test-case = { workspace = true }
 memchr = { workspace = true }
 strum = { workspace = true }
 strum_macros = { workspace = true }
+strsim = "0.11.1"
 
 [dev-dependencies]
 ruff_db = { workspace = true, features = ["testing", "os"] }

crates/ty_python_semantic/resources/mdtest/narrow/assignment.md

@@ -244,8 +244,7 @@ class D(TypedDict):
 
 td = D(x=1, label="a")
 td["x"] = 0
-# TODO: should be Literal[0]
-reveal_type(td["x"])  # revealed: @Todo(Support for `TypedDict`)
+reveal_type(td["x"])  # revealed: Literal[0]
 
 # error: [unresolved-reference]
 does["not"]["exist"] = 0

crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap

@@ -0,0 +1,77 @@
+---
+source: crates/ty_test/src/lib.rs
+expression: snapshot
+---
+---
+mdtest name: typed_dict.md - `TypedDict` - Diagnostics
+mdtest path: crates/ty_python_semantic/resources/mdtest/typed_dict.md
+---
+
+# Python source files
+
+## mdtest_snippet.py
+
+```
+ 1 | from typing import TypedDict, Final
+ 2 | 
+ 3 | class Person(TypedDict):
+ 4 |     name: str
+ 5 |     age: int | None
+ 6 | 
+ 7 | def access_invalid_literal_string_key(person: Person):
+ 8 |     person["naem"]  # error: [invalid-key]
+ 9 | 
+10 | NAME_KEY: Final = "naem"
+11 | 
+12 | def access_invalid_key(person: Person):
+13 |     person[NAME_KEY]  # error: [invalid-key]
+14 | 
+15 | def access_with_str_key(person: Person, str_key: str):
+16 |     person[str_key]  # error: [invalid-key]
+```
+
+# Diagnostics
+
+```
+error[invalid-key]: Invalid key access on TypedDict `Person`
+  --> src/mdtest_snippet.py:8:5
+   |
+ 7 | def access_invalid_literal_string_key(person: Person):
+ 8 |     person["naem"]  # error: [invalid-key]
+   |     ------ ^^^^^^ Unknown key "naem" - did you mean "name"?
+   |     |
+   |     TypedDict `Person`
+ 9 |
+10 | NAME_KEY: Final = "naem"
+   |
+info: rule `invalid-key` is enabled by default
+
+```
+
+```
+error[invalid-key]: Invalid key access on TypedDict `Person`
+  --> src/mdtest_snippet.py:13:5
+   |
+12 | def access_invalid_key(person: Person):
+13 |     person[NAME_KEY]  # error: [invalid-key]
+   |     ------ ^^^^^^^^ Unknown key "naem" - did you mean "name"?
+   |     |
+   |     TypedDict `Person`
+14 |
+15 | def access_with_str_key(person: Person, str_key: str):
+   |
+info: rule `invalid-key` is enabled by default
+
+```
+
+```
+error[invalid-key]: TypedDict `Person` cannot be indexed with a key of type `str`
+  --> src/mdtest_snippet.py:16:12
+   |
+15 | def access_with_str_key(person: Person, str_key: str):
+16 |     person[str_key]  # error: [invalid-key]
+   |            ^^^^^^^
+   |
+info: rule `invalid-key` is enabled by default
+
+```

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -25,12 +25,21 @@ alice: Person = {"name": "Alice", "age": 30}
 reveal_type(alice["name"])  # revealed: Unknown
 # TODO: this should be `int | None`
 reveal_type(alice["age"])  # revealed: Unknown
+
+# TODO: this should reveal `Unknown`, and it should emit an error
+reveal_type(alice["non_existing"])  # revealed: Unknown
 ```
 
 Inhabitants can also be created through a constructor call:
 
 ```py
 bob = Person(name="Bob", age=25)
+
+reveal_type(bob["name"])  # revealed: str
+reveal_type(bob["age"])  # revealed: int | None
+
+# error: [invalid-key] "Invalid key access on TypedDict `Person`: Unknown key "non_existing""
+reveal_type(bob["non_existing"])  # revealed: Unknown
 ```
 
 Methods that are available on `dict`s are also available on `TypedDict`s:
@@ -127,6 +136,39 @@ dangerous(alice)
 reveal_type(alice["name"])  # revealed: Unknown
 ```
 
+## Key-based access
+
+```py
+from typing import TypedDict, Final, Literal, Any
+
+class Person(TypedDict):
+    name: str
+    age: int | None
+
+NAME_FINAL: Final = "name"
+AGE_FINAL: Final[Literal["age"]] = "age"
+
+def _(person: Person, literal_key: Literal["age"], union_of_keys: Literal["age", "name"], str_key: str, unknown_key: Any) -> None:
+    reveal_type(person["name"])  # revealed: str
+    reveal_type(person["age"])  # revealed: int | None
+
+    reveal_type(person[NAME_FINAL])  # revealed: str
+    reveal_type(person[AGE_FINAL])  # revealed: int | None
+
+    reveal_type(person[literal_key])  # revealed: int | None
+
+    reveal_type(person[union_of_keys])  # revealed: int | None | str
+
+    # error: [invalid-key] "Invalid key access on TypedDict `Person`: Unknown key "non_existing""
+    reveal_type(person["non_existing"])  # revealed: Unknown
+
+    # error: [invalid-key] "TypedDict `Person` cannot be indexed with a key of type `str`"
+    reveal_type(person[str_key])  # revealed: Unknown
+
+    # No error here:
+    reveal_type(person[unknown_key])  # revealed: Unknown
+```
+
 ## Methods on `TypedDict`
 
 ```py
@@ -333,4 +375,29 @@ reveal_type(Message.__required_keys__)  # revealed: @Todo(Support for `TypedDict
 msg.content
 ```
 
+## Diagnostics
+
+<!-- snapshot-diagnostics -->
+
+Snapshot tests for diagnostic messages including suggestions:
+
+```py
+from typing import TypedDict, Final
+
+class Person(TypedDict):
+    name: str
+    age: int | None
+
+def access_invalid_literal_string_key(person: Person):
+    person["naem"]  # error: [invalid-key]
+
+NAME_KEY: Final = "naem"
+
+def access_invalid_key(person: Person):
+    person[NAME_KEY]  # error: [invalid-key]
+
+def access_with_str_key(person: Person, str_key: str):
+    person[str_key]  # error: [invalid-key]
+```
+
 [`typeddict`]: https://typing.python.org/en/latest/spec/typeddict.html

crates/ty_python_semantic/src/lib.rs

@@ -46,6 +46,7 @@ mod util;
 #[cfg(feature = "testing")]
 pub mod pull_types;
 
+type FxOrderMap<K, V> = ordermap::map::OrderMap<K, V, BuildHasherDefault<FxHasher>>;
 type FxOrderSet<V> = ordermap::set::OrderSet<V, BuildHasherDefault<FxHasher>>;
 type FxIndexMap<K, V> = indexmap::IndexMap<K, V, BuildHasherDefault<FxHasher>>;
 type FxIndexSet<V> = indexmap::IndexSet<V, BuildHasherDefault<FxHasher>>;

crates/ty_python_semantic/src/types.rs

@@ -38,6 +38,7 @@ use crate::semantic_index::scope::ScopeId;
 use crate::semantic_index::{imported_modules, place_table, semantic_index};
 use crate::suppression::check_suppressions;
 use crate::types::call::{Binding, Bindings, CallArguments, CallableBinding};
+use crate::types::class::{CodeGeneratorKind, Field};
 pub(crate) use crate::types::class_base::ClassBase;
 use crate::types::context::{LintDiagnosticGuard, LintDiagnosticGuardBuilder};
 use crate::types::diagnostic::{INVALID_TYPE_FORM, UNSUPPORTED_BOOL_CONVERSION};
@@ -61,7 +62,7 @@ use crate::types::signatures::{Parameter, ParameterForm, Parameters, walk_signat
 use crate::types::tuple::{TupleSpec, TupleType};
 use crate::unpack::EvaluationMode;
 pub use crate::util::diagnostics::add_inferred_python_version_hint_to_diagnostic;
-use crate::{Db, FxOrderSet, Module, Program};
+use crate::{Db, FxOrderMap, FxOrderSet, Module, Program};
 pub(crate) use class::{ClassLiteral, ClassType, GenericAlias, KnownClass};
 use instance::Protocol;
 pub use instance::{NominalInstanceType, ProtocolInstanceType};
@@ -669,10 +670,6 @@ impl<'db> Type<'db> {
         matches!(self, Type::Dynamic(_))
     }
 
-    pub(crate) const fn is_typed_dict(&self) -> bool {
-        matches!(self, Type::TypedDict(..))
-    }
-
     /// Returns the top materialization (or upper bound materialization) of this type, which is the
     /// most general form of the type that is fully static.
     #[must_use]
@@ -834,6 +831,17 @@ impl<'db> Type<'db> {
             .expect("Expected a Type::EnumLiteral variant")
     }
 
+    pub(crate) const fn is_typed_dict(&self) -> bool {
+        matches!(self, Type::TypedDict(..))
+    }
+
+    pub(crate) fn into_typed_dict(self) -> Option<TypedDictType<'db>> {
+        match self {
+            Type::TypedDict(typed_dict) => Some(typed_dict),
+            _ => None,
+        }
+    }
+
     /// Turn a class literal (`Type::ClassLiteral` or `Type::GenericAlias`) into a `ClassType`.
     /// Since a `ClassType` must be specialized, apply the default specialization to any
     /// unspecialized generic class literal.
@@ -5264,15 +5272,15 @@ impl<'db> Type<'db> {
                         ],
                     ),
                     _ if class.is_typed_dict(db) => {
-                        Type::TypedDict(TypedDictType::new(db, ClassType::NonGeneric(*class)))
+                        TypedDictType::from(db, ClassType::NonGeneric(*class))
                     }
                     _ => Type::instance(db, class.default_specialization(db)),
                 };
                 Ok(ty)
             }
-            Type::GenericAlias(alias) if alias.is_typed_dict(db) => Ok(Type::TypedDict(
-                TypedDictType::new(db, ClassType::from(*alias)),
-            )),
+            Type::GenericAlias(alias) if alias.is_typed_dict(db) => {
+                Ok(TypedDictType::from(db, ClassType::from(*alias)))
+            }
             Type::GenericAlias(alias) => Ok(Type::instance(db, ClassType::from(*alias))),
 
             Type::SubclassOf(_)
@@ -5619,6 +5627,7 @@ impl<'db> Type<'db> {
             return KnownClass::Dict
                 .to_specialized_class_type(db, [KnownClass::Str.to_instance(db), Type::object(db)])
                 .map(Type::from)
+                // Guard against user-customized typesheds with a broken `dict` class
                 .unwrap_or_else(Type::unknown);
         }
 
@@ -8928,6 +8937,15 @@ pub struct TypedDictType<'db> {
 impl get_size2::GetSize for TypedDictType<'_> {}
 
 impl<'db> TypedDictType<'db> {
+    pub(crate) fn from(db: &'db dyn Db, defining_class: ClassType<'db>) -> Type<'db> {
+        Type::TypedDict(Self::new(db, defining_class))
+    }
+
+    pub(crate) fn items(self, db: &'db dyn Db) -> FxOrderMap<Name, Field<'db>> {
+        let (class_literal, specialization) = self.defining_class(db).class_literal(db);
+        class_literal.fields(db, specialization, CodeGeneratorKind::TypedDict)
+    }
+
     pub(crate) fn apply_type_mapping<'a>(
         self,
         db: &'db dyn Db,