Handle attribute expressions as well

Author: sharkdp

crates/ty_python_semantic/resources/mdtest/implicit_type_aliases.md

@@ -366,7 +366,9 @@ def g(obj: Y):
     reveal_type(obj)  # revealed: list[int | str]
 ```
 
-## Generic types
+## Generic implicit type aliases
+
+### Functionality
 
 Implicit type aliases can also be generic:
 
@@ -528,6 +530,36 @@ class Derived2(GenericBaseAlias[int]):
     pass
 ```
 
+### Imported aliases
+
+Generic implicit type aliases can be imported from other modules and specialized:
+
+`my_types.py`:
+
+```py
+from typing_extensions import TypeVar
+
+T = TypeVar("T")
+
+MyList = list[T]
+```
+
+`main.py`:
+
+```py
+from my_types import MyList
+import my_types as mt
+
+def _(
+    list_of_ints1: MyList[int],
+    list_of_ints2: mt.MyList[int],
+):
+    reveal_type(list_of_ints1)  # revealed: list[int]
+    reveal_type(list_of_ints2)  # revealed: list[int]
+```
+
+### Error cases
+
 A generic alias that is already fully specialized cannot be specialized again:
 
 ```py
@@ -542,6 +574,14 @@ Specializing a generic implicit type alias with an incorrect number of type argu
 in an error:
 
 ```py
+from typing_extensions import TypeVar
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+MyList = list[T]
+MyDict = dict[T, U]
+
 def _(
     # error: [too-many-positional-arguments] "Too many positional arguments: expected 1, got 2"
     list_too_many_args: MyList[int, str],
@@ -563,12 +603,79 @@ def this_does_not_work() -> TypeOf[IntOrStr]:
     raise NotImplementedError()
 
 def _(
-    # error: [invalid-type-form] "Cannot specialize a non-name node in a type expression"
+    # error: [invalid-type-form] "Only name- and attribute expressions can be specialized in type expressions"
     specialized: this_does_not_work()[int],
 ):
     reveal_type(specialized)  # revealed: Unknown
 ```
 
+### Multiple definitions
+
+#### Shadowed definitions
+
+When a generic type alias shadows a definition from an outer scope, the inner definition is used:
+
+```py
+from typing_extensions import TypeVar
+
+T = TypeVar("T")
+
+MyAlias = list[T]
+
+def outer():
+    MyAlias = set[T]
+
+    def _(x: MyAlias[int]):
+        reveal_type(x)  # revealed: set[int]
+```
+
+#### Statically known conditions
+
+```py
+from typing_extensions import TypeVar
+
+T = TypeVar("T")
+
+if True:
+    MyAlias1 = list[T]
+else:
+    MyAlias1 = set[T]
+
+if False:
+    MyAlias2 = list[T]
+else:
+    MyAlias2 = set[T]
+
+def _(
+    x1: MyAlias1[int],
+    x2: MyAlias2[int],
+):
+    reveal_type(x1)  # revealed: list[int]
+    reveal_type(x2)  # revealed: set[int]
+```
+
+#### Statically unknown conditions
+
+If several definitions are visible, we emit an error:
+
+```py
+from typing_extensions import TypeVar
+
+T = TypeVar("T")
+
+def flag() -> bool:
+    return True
+
+if flag():
+    MyAlias = list[T]
+else:
+    MyAlias = set[T]
+
+# error: [invalid-type-form] "Invalid subscript of object of type `<class 'list[T@MyAlias]'> | <class 'set[T@MyAlias]'>` in type expression"
+def _(x: MyAlias[int]):
+    reveal_type(x)  # revealed: Unknown
+```
+
 ## `Literal`s
 
 We also support `typing.Literal` in implicit type aliases.

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

@@ -15,7 +15,7 @@ use crate::types::{
     KnownInstanceType, LintDiagnosticGuard, Parameter, Parameters, SpecialFormType, SubclassOfType,
     Type, TypeAliasType, TypeContext, TypeIsType, TypeMapping, UnionBuilder, UnionType, todo_type,
 };
-use crate::{FxOrderSet, ResolvedDefinition, definitions_for_name};
+use crate::{FxOrderSet, ResolvedDefinition, definitions_for_attribute, definitions_for_name};
 
 /// Type expressions
 impl<'db> TypeInferenceBuilder<'db, '_> {
@@ -723,18 +723,34 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
     ) -> Type<'db> {
         let db = self.db();
 
-        let Some(value) = subscript.value.as_name_expr() else {
-            if let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, subscript) {
-                builder.into_diagnostic("Cannot specialize a non-name node in a type expression");
+        let definitions = match &*subscript.value {
+            ast::Expr::Name(id) => {
+                // TODO: This is an expensive call to an API that was never meant to be called from
+                // type inference. We plan to rework how `in_type_expression` works in the future.
+                // This new approach will make this call unnecessary, so for now, we accept the hit
+                // in performance.
+                definitions_for_name(self.db(), self.file(), id)
+            }
+            ast::Expr::Attribute(attribute) => {
+                // TODO: See above
+                definitions_for_attribute(self.db(), self.file(), attribute)
+            }
+            _ => {
+                if let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, subscript) {
+                    builder.into_diagnostic(
+                    "Only name- and attribute expressions can be specialized in type expressions",
+                );
+                }
+                return Type::unknown();
             }
-            return Type::unknown();
         };
 
-        // TODO: This is an expensive call to an API that was never meant to be called from
-        // type inference. We plan to rework how `in_type_expression` works in the future.
-        // This new approach will make this call unnecessary, so for now, we accept the hit
-        // in performance.
-        let definitions = definitions_for_name(self.db(), self.file(), value);
+        // TODO: If an implicit type alias is defined multiple times, we arbitrarily pick the
+        // first definition here. Instead, we should do proper name resolution to find the
+        // definition that is actually being referenced. Similar to the comments above, this
+        // should soon be addressed by a rework of how `in_type_expression` works. In the
+        // meantime, we seem to be doing okay in practice (see "Multiple definitions" tests in
+        // `implicit_type_aliases.md`).
         let Some(type_alias_definition) =
             definitions.iter().find_map(ResolvedDefinition::definition)
         else {