Skip to content

Commit 11a9e7e

Browse files
authored
[ty] use type context more aggressively to infer values ​​when constructing a TypedDict (#20806)
## Summary Based on @ibraheemdev's comment on #20792: > I think we can also update our bidirectional inference code, [which makes the same assumption](https://github.com/astral-sh/ruff/blob/main/crates/ty_python_semantic/src/types/infer/builder.rs?rgh-link-date=2025-10-09T21%3A30%3A31Z#L5860). This PR also adds more test cases for how `TypedDict` annotations affect generic call inference. ## Test Plan New tests in `typed_dict.md`
1 parent bbd3856 commit 11a9e7e

File tree

2 files changed

+65
-4
lines changed

2 files changed

+65
-4
lines changed

crates/ty_python_semantic/resources/mdtest/typed_dict.md

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,11 @@ specific value types for each valid key. Each string key can be either required
55

66
## Basic
77

8+
```toml
9+
[environment]
10+
python-version = "3.12"
11+
```
12+
813
Here, we define a `TypedDict` using the class-based syntax:
914

1015
```py
@@ -105,6 +110,39 @@ eve3a: Person = {"name": "Eve", "age": 25, "extra": True}
105110
eve3b = Person(name="Eve", age=25, extra=True)
106111
```
107112

113+
Also, the value types ​​declared in a `TypedDict` affect generic call inference:
114+
115+
```py
116+
class Plot(TypedDict):
117+
y: list[int]
118+
x: list[int] | None
119+
120+
plot1: Plot = {"y": [1, 2, 3], "x": None}
121+
122+
def homogeneous_list[T](*args: T) -> list[T]:
123+
return list(args)
124+
125+
reveal_type(homogeneous_list(1, 2, 3)) # revealed: list[Literal[1, 2, 3]]
126+
plot2: Plot = {"y": homogeneous_list(1, 2, 3), "x": None}
127+
reveal_type(plot2["y"]) # revealed: list[int]
128+
# TODO: no error
129+
# error: [invalid-argument-type]
130+
plot3: Plot = {"y": homogeneous_list(1, 2, 3), "x": homogeneous_list(1, 2, 3)}
131+
132+
Y = "y"
133+
X = "x"
134+
135+
plot4: Plot = {Y: [1, 2, 3], X: None}
136+
plot5: Plot = {Y: homogeneous_list(1, 2, 3), X: None}
137+
138+
class Items(TypedDict):
139+
items: list[int | str]
140+
141+
items1: Items = {"items": homogeneous_list(1, 2, 3)}
142+
ITEMS = "items"
143+
items2: Items = {ITEMS: homogeneous_list(1, 2, 3)}
144+
```
145+
108146
Assignments to keys are also validated:
109147

110148
```py
@@ -796,6 +834,18 @@ p2: TaggedData[str] = {"data": "Hello", "tag": "text"}
796834

797835
# error: [invalid-argument-type] "Invalid argument to key "data" with declared type `int` on TypedDict `TaggedData`: value of type `Literal["not a number"]`"
798836
p3: TaggedData[int] = {"data": "not a number", "tag": "number"}
837+
838+
class Items(TypedDict, Generic[T]):
839+
items: list[T]
840+
841+
def homogeneous_list(*args: T) -> list[T]:
842+
return list(args)
843+
844+
items1: Items[int] = {"items": [1, 2, 3]}
845+
items2: Items[str] = {"items": ["a", "b", "c"]}
846+
items3: Items[int] = {"items": homogeneous_list(1, 2, 3)}
847+
items4: Items[str] = {"items": homogeneous_list("a", "b", "c")}
848+
items5: Items[int | str] = {"items": homogeneous_list(1, 2, 3)}
799849
```
800850

801851
### PEP-695 generics
@@ -817,6 +867,18 @@ p2: TaggedData[str] = {"data": "Hello", "tag": "text"}
817867

818868
# error: [invalid-argument-type] "Invalid argument to key "data" with declared type `int` on TypedDict `TaggedData`: value of type `Literal["not a number"]`"
819869
p3: TaggedData[int] = {"data": "not a number", "tag": "number"}
870+
871+
class Items[T](TypedDict):
872+
items: list[T]
873+
874+
def homogeneous_list[T](*args: T) -> list[T]:
875+
return list(args)
876+
877+
items1: Items[int] = {"items": [1, 2, 3]}
878+
items2: Items[str] = {"items": ["a", "b", "c"]}
879+
items3: Items[int] = {"items": homogeneous_list(1, 2, 3)}
880+
items4: Items[str] = {"items": homogeneous_list("a", "b", "c")}
881+
items5: Items[int | str] = {"items": homogeneous_list(1, 2, 3)}
820882
```
821883

822884
## Recursive `TypedDict`

crates/ty_python_semantic/src/types/infer/builder.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5858,11 +5858,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
58585858
let typed_dict_items = typed_dict.items(self.db());
58595859

58605860
for item in items {
5861-
self.infer_optional_expression(item.key.as_ref(), TypeContext::default());
5861+
let key_ty = self.infer_optional_expression(item.key.as_ref(), TypeContext::default());
58625862

5863-
if let Some(ast::Expr::StringLiteral(ref key)) = item.key
5864-
&& let Some(key) = key.as_single_part_string()
5865-
&& let Some(field) = typed_dict_items.get(key.as_str())
5863+
if let Some(Type::StringLiteral(key)) = key_ty
5864+
&& let Some(field) = typed_dict_items.get(key.value(self.db()))
58665865
{
58675866
self.infer_expression(&item.value, TypeContext::new(Some(field.declared_ty)));
58685867
} else {

0 commit comments

Comments
 (0)