[ty] tighten up handling of subscripts in type expressions

Author: carljm

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

@@ -98,7 +98,9 @@ async def outer_async():  # avoid unrelated syntax errors on `yield` and `await`
         n: 1 < 2,  # error: [invalid-type-form] "Comparison expressions are not allowed in type expressions"
         o: bar(),  # error: [invalid-type-form] "Function calls are not allowed in type expressions"
         p: int | f"foo",  # error: [invalid-type-form] "F-strings are not allowed in type expressions"
-        q: [1, 2, 3][1:2],  # error: [invalid-type-form] "Slices are not allowed in type expressions"
+        # error: [invalid-type-form] "Slices are not allowed in type expressions"
+        # error: [invalid-type-form] "Cannot subscript object of type"
+        q: [1, 2, 3][1:2],
     ):
         reveal_type(a)  # revealed: Unknown
         reveal_type(b)  # revealed: Unknown
@@ -116,7 +118,7 @@ async def outer_async():  # avoid unrelated syntax errors on `yield` and `await`
         reveal_type(n)  # revealed: Unknown
         reveal_type(o)  # revealed: Unknown
         reveal_type(p)  # revealed: int | Unknown
-        reveal_type(q)  # revealed: @Todo(unknown type subscript)
+        reveal_type(q)  # revealed: Unknown
 
 class Mat:
     def __init__(self, value: int):

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

@@ -330,10 +330,11 @@ from other import Literal
 # ?
 #
 # error: [invalid-type-form] "Int literals are not allowed in this context in a type expression"
+# error: [invalid-type-form] "Cannot subscript object of type `_SpecialForm` in type expression"
 a1: Literal[26]
 
 def f():
-    reveal_type(a1)  # revealed: @Todo(unknown type subscript)
+    reveal_type(a1)  # revealed: Unknown
 ```
 
 ## Detecting typing_extensions.Literal

crates/ty_python_semantic/resources/mdtest/implicit_type_aliases.md

@@ -33,9 +33,11 @@ g(None)
 We also support unions in type aliases:
 
 ```py
-from typing_extensions import Any, Never, Literal, LiteralString, Tuple, Annotated, Optional, Union
+from typing_extensions import Any, Never, Literal, LiteralString, Tuple, Annotated, Optional, Union, TypeVar
 from ty_extensions import Unknown
 
+T = TypeVar("T")
+
 IntOrStr = int | str
 IntOrStrOrBytes1 = int | str | bytes
 IntOrStrOrBytes2 = (int | str) | bytes
@@ -68,6 +70,10 @@ IntOrOptional = int | Optional[str]
 OptionalOrInt = Optional[str] | int
 IntOrTypeOfStr = int | type[str]
 TypeOfStrOrInt = type[str] | int
+TypeVarOrInt = T | int
+IntOrTypeVar = int | T
+TypeVarOrNone = T | None
+NoneOrTypeVar = None | T
 
 reveal_type(IntOrStr)  # revealed: types.UnionType
 reveal_type(IntOrStrOrBytes1)  # revealed: types.UnionType
@@ -101,6 +107,10 @@ reveal_type(IntOrOptional)  # revealed: types.UnionType
 reveal_type(OptionalOrInt)  # revealed: types.UnionType
 reveal_type(IntOrTypeOfStr)  # revealed: types.UnionType
 reveal_type(TypeOfStrOrInt)  # revealed: types.UnionType
+reveal_type(TypeVarOrInt)  # revealed: types.UnionType
+reveal_type(IntOrTypeVar)  # revealed: types.UnionType
+reveal_type(TypeVarOrNone)  # revealed: types.UnionType
+reveal_type(NoneOrTypeVar)  # revealed: types.UnionType
 
 def _(
     int_or_str: IntOrStr,
@@ -135,6 +145,10 @@ def _(
     optional_or_int: OptionalOrInt,
     int_or_type_of_str: IntOrTypeOfStr,
     type_of_str_or_int: TypeOfStrOrInt,
+    type_var_or_int: TypeVarOrInt,
+    int_or_type_var: IntOrTypeVar,
+    type_var_or_none: TypeVarOrNone,
+    none_or_type_var: NoneOrTypeVar,
 ):
     reveal_type(int_or_str)  # revealed: int | str
     reveal_type(int_or_str_or_bytes1)  # revealed: int | str | bytes
@@ -168,6 +182,14 @@ def _(
     reveal_type(optional_or_int)  # revealed: str | None | int
     reveal_type(int_or_type_of_str)  # revealed: int | type[str]
     reveal_type(type_of_str_or_int)  # revealed: type[str] | int
+    # TODO should be Unknown | int
+    reveal_type(type_var_or_int)  # revealed: T@_ | int
+    # TODO should be int | Unknown
+    reveal_type(int_or_type_var)  # revealed: int | T@_
+    # TODO should be Unknown | None
+    reveal_type(type_var_or_none)  # revealed: T@_ | None
+    # TODO should be None | Unknown
+    reveal_type(none_or_type_var)  # revealed: None | T@_
 ```
 
 If a type is unioned with itself in a value expression, the result is just that type. No
@@ -349,7 +371,7 @@ MyList = list[T]
 
 def _(my_list: MyList[int]):
     # TODO: This should be `list[int]`
-    reveal_type(my_list)  # revealed: @Todo(unknown type subscript)
+    reveal_type(my_list)  # revealed: @Todo(specialized generic alias in type expression)
 
 ListOrTuple = list[T] | tuple[T, ...]
 

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

@@ -9506,6 +9506,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     KnownInstanceType::UnionType(_)
                     | KnownInstanceType::Literal(_)
                     | KnownInstanceType::Annotated(_)
+                    | KnownInstanceType::TypeVar(_)
                     | KnownInstanceType::TypeGenericAlias(_),
                 ),
                 Type::ClassLiteral(..)
@@ -9516,6 +9517,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     KnownInstanceType::UnionType(_)
                     | KnownInstanceType::Literal(_)
                     | KnownInstanceType::Annotated(_)
+                    | KnownInstanceType::TypeVar(_)
                     | KnownInstanceType::TypeGenericAlias(_),
                 ),
                 ast::Operator::BitOr,
@@ -10916,6 +10918,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     .map(Type::from)
                     .unwrap_or_else(Type::unknown);
             }
+            Type::KnownInstance(KnownInstanceType::UnionType(_)) => {
+                return todo_type!("Specialization of union type alias");
+            }
             _ => {}
         }
 

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

@@ -854,7 +854,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                     Type::unknown()
                 }
             },
-            Type::Dynamic(DynamicType::Todo(_)) => {
+            Type::Dynamic(_) => {
                 self.infer_type_expression(slice);
                 value_ty
             }
@@ -883,11 +883,22 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                     }
                 }
             }
+            Type::GenericAlias(_) => {
+                self.infer_type_expression(slice);
+                // If the generic alias is already fully specialized, this is an error. But it
+                // could have been specialized with another typevar (e.g. a type alias like `MyList
+                // = list[T]`), in which case it's later valid to do `MyList[int]`.
+                todo_type!("specialized generic alias in type expression")
+            }
             _ => {
-                // TODO: Emit a diagnostic once we've implemented all valid subscript type
-                // expressions.
                 self.infer_type_expression(slice);
-                todo_type!("unknown type subscript")
+                if let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, subscript) {
+                    builder.into_diagnostic(format_args!(
+                        "Cannot subscript object of type `{}` in type expression",
+                        value_ty.display(self.db())
+                    ));
+                }
+                Type::unknown()
             }
         }
     }