[ty] Infer slightly more precise types for comprehensions (#20111)

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

@@ -48,6 +48,24 @@ def _(
         reveal_type(h_)  # revealed: Unknown
         reveal_type(i_)  # revealed: Unknown
         reveal_type(j_)  # revealed: Unknown
+
+# Inspired by the conformance test suite at
+# https://github.com/python/typing/blob/d4f39b27a4a47aac8b6d4019e1b0b5b3156fabdc/conformance/tests/aliases_implicit.py#L88-L122
+B = [x for x in range(42)]
+C = {x for x in range(42)}
+D = {x: y for x, y in enumerate(range(42))}
+E = (x for x in range(42))
+
+def _(
+    b: B,  # error: [invalid-type-form]
+    c: C,  # error: [invalid-type-form]
+    d: D,  # error: [invalid-type-form]
+    e: E,  # error: [invalid-type-form]
+):
+    reveal_type(b)  # revealed: Unknown
+    reveal_type(c)  # revealed: Unknown
+    reveal_type(d)  # revealed: Unknown
+    reveal_type(e)  # revealed: Unknown
 ```
 
 ## Invalid AST nodes

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

@@ -75,7 +75,7 @@ a: tuple[()] = (1, 2)
 # error: [invalid-assignment] "Object of type `tuple[Literal["foo"]]` is not assignable to `tuple[int]`"
 b: tuple[int] = ("foo",)
 
-# error: [invalid-assignment] "Object of type `tuple[list[Unknown], Literal["foo"]]` is not assignable to `tuple[str | int, str]`"
+# error: [invalid-assignment]
 c: tuple[str | int, str] = ([], "foo")
 ```
 

crates/ty_python_semantic/resources/mdtest/del.md

@@ -46,7 +46,7 @@ def delete():
     del d  # error: [unresolved-reference] "Name `d` used when not defined"
 
 delete()
-reveal_type(d)  # revealed: list[Unknown]
+reveal_type(d)  # revealed: list[@Todo(list literal element type)]
 
 def delete_element():
     # When the `del` target isn't a name, it doesn't force local resolution.
@@ -62,7 +62,7 @@ def delete_global():
 
 delete_global()
 # Again, the variable should have been removed, but we don't check it.
-reveal_type(d)  # revealed: list[Unknown]
+reveal_type(d)  # revealed: list[@Todo(list literal element type)]
 
 def delete_nonlocal():
     e = 2

crates/ty_python_semantic/resources/mdtest/directives/assert_never.md

@@ -33,12 +33,6 @@ def _():
 def _():
     assert_never(None)  # error: [type-assertion-failure]
 
-def _():
-    assert_never([])  # error: [type-assertion-failure]
-
-def _():
-    assert_never({})  # error: [type-assertion-failure]
-
 def _():
     assert_never(())  # error: [type-assertion-failure]
 

crates/ty_python_semantic/resources/mdtest/import/dunder_all.md

@@ -785,7 +785,7 @@ from subexporter import *
 
 # TODO: Should be `list[str]`
 # TODO: Should we avoid including `Unknown` for this case?
-reveal_type(__all__)  # revealed: Unknown | list[Unknown]
+reveal_type(__all__)  # revealed: Unknown | list[@Todo(list literal element type)]
 
 __all__.append("B")
 

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

@@ -3,5 +3,12 @@
 ## Empty dictionary
 
 ```py
-reveal_type({})  # revealed: dict[Unknown, Unknown]
+reveal_type({})  # revealed: dict[@Todo(dict literal key type), @Todo(dict literal value type)]
+```
+
+## Dict comprehensions
+
+```py
+# revealed: dict[@Todo(dict comprehension key type), @Todo(dict comprehension value type)]
+reveal_type({x: y for x, y in enumerate(range(42))})
 ```

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

@@ -0,0 +1,6 @@
+# Generator expressions
+
+```py
+# revealed: GeneratorType[@Todo(generator expression yield type), @Todo(generator expression send type), @Todo(generator expression return type)]
+reveal_type((x for x in range(42)))
+```

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

@@ -3,5 +3,11 @@
 ## Empty list
 
 ```py
-reveal_type([])  # revealed: list[Unknown]
+reveal_type([])  # revealed: list[@Todo(list literal element type)]
+```
+
+## List comprehensions
+
+```py
+reveal_type([x for x in range(42)])  # revealed: list[@Todo(list comprehension element type)]
 ```

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

@@ -3,5 +3,11 @@
 ## Basic set
 
 ```py
-reveal_type({1, 2})  # revealed: set[Unknown]
+reveal_type({1, 2})  # revealed: set[@Todo(set literal element type)]
+```
+
+## Set comprehensions
+
+```py
+reveal_type({x for x in range(42)})  # revealed: set[@Todo(set comprehension element type)]
 ```

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

@@ -207,7 +207,7 @@ 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[Unknown, Unknown]
+reveal_type(cm)  # revealed: ChainMap[@Todo(dict literal key type), @Todo(dict literal value type)]
 
 reveal_type(l[0])  # revealed: Literal[0]
 reveal_type(d[0])  # revealed: Literal[0]

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/nested.md

@@ -318,7 +318,7 @@ def f(l: list[str | None]):
         l: list[str | None] = [None]
         def _():
             # TODO: should be `str | None`
-            reveal_type(l[0])  # revealed: Unknown
+            reveal_type(l[0])  # revealed: @Todo(list literal element type)
 
     def _():
         def _():

crates/ty_python_semantic/resources/mdtest/snapshots/assert_never.md_-_`assert_never`_-_Basic_functionality_-_Diagnostics_(be8f5d8b0718ee54).snap

@@ -25,22 +25,16 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/directives/assert_never.
 11 |     assert_never(None)  # error: [type-assertion-failure]
 12 | 
 13 | def _():
-14 |     assert_never([])  # error: [type-assertion-failure]
+14 |     assert_never(())  # error: [type-assertion-failure]
 15 | 
-16 | def _():
+16 | def _(flag: bool, never: Never):
-17 |     assert_never({})  # error: [type-assertion-failure]
+17 |     assert_never(1 if flag else never)  # error: [type-assertion-failure]
 18 | 
-19 | def _():
+19 | def _(any_: Any):
-20 |     assert_never(())  # error: [type-assertion-failure]
+20 |     assert_never(any_)  # error: [type-assertion-failure]
 21 | 
-22 | def _(flag: bool, never: Never):
+22 | def _(unknown: Unknown):
-23 |     assert_never(1 if flag else never)  # error: [type-assertion-failure]
+23 |     assert_never(unknown)  # error: [type-assertion-failure]
-24 | 
-25 | def _(any_: Any):
-26 |     assert_never(any_)  # error: [type-assertion-failure]
-27 | 
-28 | def _(unknown: Unknown):
-29 |     assert_never(unknown)  # error: [type-assertion-failure]
 ```
 
 # Diagnostics
@@ -101,46 +95,12 @@ error[type-assertion-failure]: Argument does not have asserted type `Never`
   --> src/mdtest_snippet.py:14:5
    |
 13 | def _():
-14 |     assert_never([])  # error: [type-assertion-failure]
+14 |     assert_never(())  # error: [type-assertion-failure]
-   |     ^^^^^^^^^^^^^--^
-   |                  |
-   |                  Inferred type of argument is `list[Unknown]`
-15 |
-16 | def _():
-   |
-info: `Never` and `list[Unknown]` are not equivalent types
-info: rule `type-assertion-failure` is enabled by default
-
-```
-
-```
-error[type-assertion-failure]: Argument does not have asserted type `Never`
-  --> src/mdtest_snippet.py:17:5
-   |
-16 | def _():
-17 |     assert_never({})  # error: [type-assertion-failure]
-   |     ^^^^^^^^^^^^^--^
-   |                  |
-   |                  Inferred type of argument is `dict[Unknown, Unknown]`
-18 |
-19 | def _():
-   |
-info: `Never` and `dict[Unknown, Unknown]` are not equivalent types
-info: rule `type-assertion-failure` is enabled by default
-
-```
-
-```
-error[type-assertion-failure]: Argument does not have asserted type `Never`
-  --> src/mdtest_snippet.py:20:5
-   |
-19 | def _():
-20 |     assert_never(())  # error: [type-assertion-failure]
    |     ^^^^^^^^^^^^^--^
    |                  |
    |                  Inferred type of argument is `tuple[()]`
-21 |
+15 |
-22 | def _(flag: bool, never: Never):
+16 | def _(flag: bool, never: Never):
    |
 info: `Never` and `tuple[()]` are not equivalent types
 info: rule `type-assertion-failure` is enabled by default
@@ -149,15 +109,15 @@ info: rule `type-assertion-failure` is enabled by default
 
 ```
 error[type-assertion-failure]: Argument does not have asserted type `Never`
-  --> src/mdtest_snippet.py:23:5
+  --> src/mdtest_snippet.py:17:5
    |
-22 | def _(flag: bool, never: Never):
+16 | def _(flag: bool, never: Never):
-23 |     assert_never(1 if flag else never)  # error: [type-assertion-failure]
+17 |     assert_never(1 if flag else never)  # error: [type-assertion-failure]
    |     ^^^^^^^^^^^^^--------------------^
    |                  |
    |                  Inferred type of argument is `Literal[1]`
-24 |
+18 |
-25 | def _(any_: Any):
+19 | def _(any_: Any):
    |
 info: `Never` and `Literal[1]` are not equivalent types
 info: rule `type-assertion-failure` is enabled by default
@@ -166,15 +126,15 @@ info: rule `type-assertion-failure` is enabled by default
 
 ```
 error[type-assertion-failure]: Argument does not have asserted type `Never`
-  --> src/mdtest_snippet.py:26:5
+  --> src/mdtest_snippet.py:20:5
    |
-25 | def _(any_: Any):
+19 | def _(any_: Any):
-26 |     assert_never(any_)  # error: [type-assertion-failure]
+20 |     assert_never(any_)  # error: [type-assertion-failure]
    |     ^^^^^^^^^^^^^----^
    |                  |
    |                  Inferred type of argument is `Any`
-27 |
+21 |
-28 | def _(unknown: Unknown):
+22 | def _(unknown: Unknown):
    |
 info: `Never` and `Any` are not equivalent types
 info: rule `type-assertion-failure` is enabled by default
@@ -183,10 +143,10 @@ info: rule `type-assertion-failure` is enabled by default
 
 ```
 error[type-assertion-failure]: Argument does not have asserted type `Never`
-  --> src/mdtest_snippet.py:29:5
+  --> src/mdtest_snippet.py:23:5
    |
-28 | def _(unknown: Unknown):
+22 | def _(unknown: Unknown):
-29 |     assert_never(unknown)  # error: [type-assertion-failure]
+23 |     assert_never(unknown)  # error: [type-assertion-failure]
    |     ^^^^^^^^^^^^^-------^
    |                  |
    |                  Inferred type of argument is `Unknown`

crates/ty_python_semantic/resources/mdtest/subscript/lists.md

@@ -9,13 +9,13 @@ A list can be indexed into with:
 
 ```py
 x = [1, 2, 3]
-reveal_type(x)  # revealed: list[Unknown]
+reveal_type(x)  # revealed: list[@Todo(list literal element type)]
 
 # TODO reveal int
-reveal_type(x[0])  # revealed: Unknown
+reveal_type(x[0])  # revealed: @Todo(list literal element type)
 
 # TODO reveal list[int]
-reveal_type(x[0:1])  # revealed: list[Unknown]
+reveal_type(x[0:1])  # revealed: list[@Todo(list literal element type)]
 
 # error: [invalid-argument-type]
 reveal_type(x["a"])  # revealed: Unknown

crates/ty_python_semantic/resources/mdtest/type_compendium/tuple.md

@@ -47,9 +47,9 @@ def f(x: Iterable[int], y: list[str], z: Never, aa: list[Never]):
 reveal_type(tuple((1, 2)))  # revealed: tuple[Literal[1], Literal[2]]
 
 # TODO: should be `tuple[Literal[1], ...]`
-reveal_type(tuple([1]))  # revealed: tuple[Unknown, ...]
+reveal_type(tuple([1]))  # revealed: tuple[@Todo(list literal element type), ...]
 
-# error: [invalid-argument-type] "Argument is incorrect: Expected `tuple[int]`, found `list[Unknown]`"
+# error: [invalid-argument-type]
 reveal_type(tuple[int]([1]))  # revealed: tuple[int]
 
 # error: [invalid-argument-type] "Argument is incorrect: Expected `tuple[int, str]`, found `tuple[Literal[1]]`"

crates/ty_python_semantic/resources/mdtest/unpacking.md

@@ -214,8 +214,8 @@ reveal_type(d)  # revealed: Literal[2]
 ```py
 a, b = [1, 2]
 # TODO: should be `int` for both `a` and `b`
-reveal_type(a)  # revealed: Unknown
+reveal_type(a)  # revealed: @Todo(list literal element type)
-reveal_type(b)  # revealed: Unknown
+reveal_type(b)  # revealed: @Todo(list literal element type)
 ```
 
 ### Simple unpacking

crates/ty_python_semantic/src/types/infer.rs

@@ -5805,8 +5805,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             self.infer_expression(elt);
         }
 
-        // TODO generic
+        KnownClass::List
-        KnownClass::List.to_instance(self.db())
+            .to_specialized_instance(self.db(), [todo_type!("list literal element type")])
     }
 
     fn infer_set_expression(&mut self, set: &ast::ExprSet) -> Type<'db> {
@@ -5820,8 +5820,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             self.infer_expression(elt);
         }
 
-        // TODO generic
+        KnownClass::Set.to_specialized_instance(self.db(), [todo_type!("set literal element type")])
-        KnownClass::Set.to_instance(self.db())
     }
 
     fn infer_dict_expression(&mut self, dict: &ast::ExprDict) -> Type<'db> {
@@ -5836,8 +5835,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             self.infer_expression(&item.value);
         }
 
-        // TODO generic
+        KnownClass::Dict.to_specialized_instance(
-        KnownClass::Dict.to_instance(self.db())
+            self.db(),
+            [
+                todo_type!("dict literal key type"),
+                todo_type!("dict literal value type"),
+            ],
+        )
     }
 
     /// Infer the type of the `iter` expression of the first comprehension.
@@ -5860,7 +5864,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         self.infer_first_comprehension_iter(generators);
 
-        todo_type!("generator type")
+        KnownClass::GeneratorType.to_specialized_instance(
+            self.db(),
+            [
+                todo_type!("generator expression yield type"),
+                todo_type!("generator expression send type"),
+                todo_type!("generator expression return type"),
+            ],
+        )
     }
 
     fn infer_list_comprehension_expression(&mut self, listcomp: &ast::ExprListComp) -> Type<'db> {
@@ -5873,7 +5884,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         self.infer_first_comprehension_iter(generators);
 
-        todo_type!("list comprehension type")
+        KnownClass::List
+            .to_specialized_instance(self.db(), [todo_type!("list comprehension element type")])
     }
 
     fn infer_dict_comprehension_expression(&mut self, dictcomp: &ast::ExprDictComp) -> Type<'db> {
@@ -5887,7 +5899,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         self.infer_first_comprehension_iter(generators);
 
-        todo_type!("dict comprehension type")
+        KnownClass::Dict.to_specialized_instance(
+            self.db(),
+            [
+                todo_type!("dict comprehension key type"),
+                todo_type!("dict comprehension value type"),
+            ],
+        )
     }
 
     fn infer_set_comprehension_expression(&mut self, setcomp: &ast::ExprSetComp) -> Type<'db> {
@@ -5900,7 +5918,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         self.infer_first_comprehension_iter(generators);
 
-        todo_type!("set comprehension type")
+        KnownClass::Set
+            .to_specialized_instance(self.db(), [todo_type!("set comprehension element type")])
     }
 
     fn infer_generator_expression_scope(&mut self, generator: &ast::ExprGenerator) {