[ty] More precise type inference for dictionary literals (#20523)

## Summary

Extends #20360 to dictionary
literals. This also improves our `TypeDict` support by passing through
nested type context.

Author: ibraheemdev

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

@@ -139,6 +139,15 @@ reveal_type(n)  # revealed: list[Literal[1, 2, 3]]
 # error: [invalid-assignment] "Object of type `list[Unknown | str]` is not assignable to `list[LiteralString]`"
 o: list[typing.LiteralString] = ["a", "b", "c"]
 reveal_type(o)  # revealed: list[LiteralString]
+
+p: dict[int, int] = {}
+reveal_type(p)  # revealed: dict[int, int]
+
+q: dict[int | str, int] = {1: 1, 2: 2, 3: 3}
+reveal_type(q)  # revealed: dict[int | str, int]
+
+r: dict[int | str, int | str] = {1: 1, 2: 2, 3: 3}
+reveal_type(r)  # revealed: dict[int | str, int | str]
 ```
 
 ## Incorrect collection literal assignments are complained aobut

crates/ty_python_semantic/resources/mdtest/call/builtins.md

@@ -57,7 +57,7 @@ type("Foo", Base, {})
 # error: [invalid-argument-type] "Argument to class `type` is incorrect: Expected `tuple[type, ...]`, found `tuple[Literal[1], Literal[2]]`"
 type("Foo", (1, 2), {})
 
-# TODO: this should be an error
+# error: [invalid-argument-type] "Argument to class `type` is incorrect: Expected `dict[str, Any]`, found `dict[Unknown | bytes, Unknown | int]`"
 type("Foo", (Base,), {b"attr": 1})
 ```
 

crates/ty_python_semantic/resources/mdtest/literal/collections/dictionary.md

@@ -3,7 +3,49 @@
 ## Empty dictionary
 
 ```py
-reveal_type({})  # revealed: dict[@Todo(dict literal key type), @Todo(dict literal value type)]
+reveal_type({})  # revealed: dict[Unknown, Unknown]
+```
+
+## Basic dict
+
+```py
+reveal_type({1: 1, 2: 1})  # revealed: dict[Unknown | int, Unknown | int]
+```
+
+## Dict of tuples
+
+```py
+reveal_type({1: (1, 2), 2: (3, 4)})  # revealed: dict[Unknown | int, Unknown | tuple[int, int]]
+```
+
+## Unpacked dict
+
+```py
+a = {"a": 1, "b": 2}
+b = {"c": 3, "d": 4}
+
+d = {**a, **b}
+reveal_type(d)  # revealed: dict[Unknown | str, Unknown | int]
+```
+
+## Dict of functions
+
+```py
+def a(_: int) -> int:
+    return 0
+
+def b(_: int) -> int:
+    return 1
+
+x = {1: a, 2: b}
+reveal_type(x)  # revealed: dict[Unknown | int, Unknown | ((_: int) -> int)]
+```
+
+## Mixed dict
+
+```py
+# revealed: dict[Unknown | str, Unknown | int | tuple[int, int] | tuple[int, int, int]]
+reveal_type({"a": 1, "b": (1, 2), "c": (1, 2, 3)})
 ```
 
 ## Dict comprehensions

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

@@ -206,8 +206,7 @@ dd: defaultdict[int, int] = defaultdict(int)
 dd[0] = 0
 cm: ChainMap[int, int] = ChainMap({1: 1}, {0: 0})
 cm[0] = 0
-# TODO: should be ChainMap[int, int]
-reveal_type(cm)  # revealed: ChainMap[@Todo(dict literal key type), @Todo(dict literal value type)]
+reveal_type(cm)  # revealed: ChainMap[Unknown | int, Unknown | int]
 
 reveal_type(l[0])  # revealed: Literal[0]
 reveal_type(d[0])  # revealed: Literal[0]

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -85,6 +85,34 @@ alice["extra"] = True
 bob["extra"] = True
 ```
 
+## Nested `TypedDict`
+
+Nested `TypedDict` fields are also supported.
+
+```py
+from typing import TypedDict
+
+class Inner(TypedDict):
+    name: str
+    age: int | None
+
+class Person(TypedDict):
+    inner: Inner
+```
+
+```py
+alice: Person = {"inner": {"name": "Alice", "age": 30}}
+
+reveal_type(alice["inner"]["name"])  # revealed: str
+reveal_type(alice["inner"]["age"])  # revealed: int | None
+
+# error: [invalid-key] "Invalid key access on TypedDict `Inner`: Unknown key "non_existing""
+reveal_type(alice["inner"]["non_existing"])  # revealed: Unknown
+
+# error: [invalid-key] "Invalid key access on TypedDict `Inner`: Unknown key "extra""
+alice: Person = {"inner": {"name": "Alice", "age": 30, "extra": 1}}
+```
+
 ## Validation of `TypedDict` construction
 
 ```py

crates/ty_python_semantic/src/types.rs

@@ -849,6 +849,28 @@ impl<'db> Type<'db> {
         matches!(self, Type::Dynamic(_))
     }
 
+    // If the type is a specialized instance of the given `KnownClass`, returns the specialization.
+    pub(crate) fn known_specialization(
+        self,
+        known_class: KnownClass,
+        db: &'db dyn Db,
+    ) -> Option<Specialization<'db>> {
+        let class_type = match self {
+            Type::NominalInstance(instance) => instance,
+            Type::TypeAlias(alias) => alias.value_type(db).into_nominal_instance()?,
+            _ => return None,
+        }
+        .class(db);
+
+        if !class_type.is_known(db, known_class) {
+            return None;
+        }
+
+        class_type
+            .into_generic_alias()
+            .map(|generic_alias| generic_alias.specialization(db))
+    }
+
     /// 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]

crates/ty_python_semantic/src/types/infer.rs

@@ -386,20 +386,8 @@ impl<'db> TypeContext<'db> {
         known_class: KnownClass,
         db: &'db dyn Db,
     ) -> Option<Specialization<'db>> {
-        let class_type = match self.annotation? {
-            Type::NominalInstance(instance) => instance,
-            Type::TypeAlias(alias) => alias.value_type(db).into_nominal_instance()?,
-            _ => return None,
-        }
-        .class(db);
-
-        if !class_type.is_known(db, known_class) {
-            return None;
-        }
-
-        class_type
-            .into_generic_alias()
-            .map(|generic_alias| generic_alias.specialization(db))
+        self.annotation
+            .and_then(|ty| ty.known_specialization(known_class, db))
     }
 }