[ty] Implicit type aliases: support for PEP 604 unions

Author: sharkdp

crates/ruff_benchmark/benches/ty_walltime.rs

@@ -146,7 +146,7 @@ static FREQTRADE: Benchmark = Benchmark::new(
         max_dep_date: "2025-06-17",
         python_version: PythonVersion::PY312,
     },
-    400,
+    500,
 );
 
 static PANDAS: Benchmark = Benchmark::new(

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

@@ -72,9 +72,6 @@ def f(x: Union) -> None:
 
 ## Implicit type aliases using new-style unions
 
-We don't recognize these as type aliases yet, but we also don't emit false-positive diagnostics if
-you use them in type expressions:
-
 ```toml
 [environment]
 python-version = "3.10"
@@ -84,5 +81,5 @@ python-version = "3.10"
 X = int | str
 
 def f(y: X):
-    reveal_type(y)  # revealed: @Todo(Support for `types.UnionType` instances in type expressions)
+    reveal_type(y)  # revealed: int | str
 ```

crates/ty_python_semantic/resources/mdtest/implicit_type_aliases.md

@@ -17,6 +17,175 @@ def f(x: MyInt):
 f(1)
 ```
 
+## None
+
+```py
+MyNone = None
+
+# TODO: this should not be an error
+# error: [invalid-type-form] "Variable of type `None` is not allowed in a type expression"
+def g(x: MyNone):
+    # TODO: this should be `None`
+    reveal_type(x)  # revealed: Unknown
+
+g(None)
+```
+
+## Unions
+
+We also support unions in type aliases:
+
+```py
+from typing_extensions import Any, Never
+from ty_extensions import Unknown
+
+IntOrStr = int | str
+IntOrStrOrBytes1 = int | str | bytes
+IntOrStrOrBytes2 = (int | str) | bytes
+IntOrStrOrBytes3 = int | (str | bytes)
+IntOrStrOrBytes4 = IntOrStr | bytes
+BytesOrIntOrStr = bytes | IntOrStr
+IntOrNone = int | None
+NoneOrInt = None | int
+IntOrStrOrNone = IntOrStr | None
+NoneOrIntOrStr = None | IntOrStr
+IntOrAny = int | Any
+AnyOrInt = Any | int
+NoneOrAny = None | Any
+AnyOrNone = Any | None
+NeverOrAny = Never | Any
+AnyOrNever = Any | Never
+UnknownOrInt = Unknown | int
+IntOrUnknown = int | Unknown
+
+def _(
+    int_or_str: IntOrStr,
+    int_or_str_or_bytes1: IntOrStrOrBytes1,
+    int_or_str_or_bytes2: IntOrStrOrBytes2,
+    int_or_str_or_bytes3: IntOrStrOrBytes3,
+    int_or_str_or_bytes4: IntOrStrOrBytes4,
+    bytes_or_int_or_str: BytesOrIntOrStr,
+    int_or_none: IntOrNone,
+    none_or_int: NoneOrInt,
+    int_or_str_or_none: IntOrStrOrNone,
+    none_or_int_or_str: NoneOrIntOrStr,
+    int_or_any: IntOrAny,
+    any_or_int: AnyOrInt,
+    none_or_any: NoneOrAny,
+    any_or_none: AnyOrNone,
+    never_or_any: NeverOrAny,
+    any_or_never: AnyOrNever,
+    unknown_or_int: UnknownOrInt,
+    int_or_unknown: IntOrUnknown,
+):
+    reveal_type(int_or_str)  # revealed: int | str
+    reveal_type(int_or_str_or_bytes1)  # revealed: int | str | bytes
+    reveal_type(int_or_str_or_bytes2)  # revealed: int | str | bytes
+    reveal_type(int_or_str_or_bytes3)  # revealed: int | str | bytes
+    reveal_type(int_or_str_or_bytes4)  # revealed: int | str | bytes
+    reveal_type(bytes_or_int_or_str)  # revealed: bytes | int | str
+    reveal_type(int_or_none)  # revealed: int | None
+    reveal_type(none_or_int)  # revealed: None | int
+    reveal_type(int_or_str_or_none)  # revealed: int | str | None
+    reveal_type(none_or_int_or_str)  # revealed: None | int | str
+    reveal_type(int_or_any)  # revealed: int | Any
+    reveal_type(any_or_int)  # revealed: Any | int
+    reveal_type(none_or_any)  # revealed: None | Any
+    reveal_type(any_or_none)  # revealed: Any | None
+    reveal_type(never_or_any)  # revealed: Any
+    reveal_type(any_or_never)  # revealed: Any
+    reveal_type(unknown_or_int)  # revealed: Unknown | int
+    reveal_type(int_or_unknown)  # revealed: int | Unknown
+```
+
+If a type is unioned with itself in a value expression, the result is just that type. No
+`types.UnionType` instance is created:
+
+```py
+IntOrInt = int | int
+ListOfIntOrListOfInt = list[int] | list[int]
+
+reveal_type(IntOrInt)  # revealed: <class 'int'>
+reveal_type(ListOfIntOrListOfInt)  # revealed: <class 'list[int]'>
+
+def _(int_or_int: IntOrInt, list_of_int_or_list_of_int: ListOfIntOrListOfInt):
+    reveal_type(int_or_int)  # revealed: int
+    reveal_type(list_of_int_or_list_of_int)  # revealed: list[int]
+```
+
+`NoneType` has no special or-operator behavior, so this is an error:
+
+```py
+None | None  # error: [unsupported-operator] "Operator `|` is unsupported between objects of type `None` and `None`"
+```
+
+When constructing something non-sensical like `int | 1`, we could ideally emit a diagnostic for the
+expression itself, as it leads to a `TypeError` at runtime. No other type checker supports this, so
+for now we only emit an error when it is used in a type expression:
+
+```py
+IntOrOne = int | 1
+
+# error: [invalid-type-form] "Variable of type `Literal[1]` is not allowed in a type expression"
+def _(int_or_one: IntOrOne):
+    reveal_type(int_or_one)  # revealed: Unknown
+```
+
+## Generic types
+
+Implicit type aliases can also refer to generic types:
+
+```py
+from typing_extensions import TypeVar
+
+T = TypeVar("T")
+
+MyList = list[T]
+
+def _(my_list: MyList[int]):
+    # TODO: This should be `list[int]`
+    reveal_type(my_list)  # revealed: @Todo(unknown type subscript)
+
+ListOrTuple = list[T] | tuple[T, ...]
+
+def _(list_or_tuple: ListOrTuple[int]):
+    reveal_type(list_or_tuple)  # revealed: @Todo(Generic specialization of types.UnionType)
+```
+
+## Stringified annotations?
+
+From the [typing spec on type aliases](https://typing.python.org/en/latest/spec/aliases.html):
+
+> Type aliases may be as complex as type hints in annotations – anything that is acceptable as a
+> type hint is acceptable in a type alias
+
+However, no other type checker seems to support stringified annotations in implicit type aliases. We
+currently also do not support them:
+
+```py
+AliasForStr = "str"
+
+# error: [invalid-type-form] "Variable of type `Literal["str"]` is not allowed in a type expression"
+def _(s: AliasForStr):
+    reveal_type(s)  # revealed: Unknown
+
+IntOrStr = int | "str"
+
+# error: [invalid-type-form] "Variable of type `Literal["str"]` is not allowed in a type expression"
+def _(int_or_str: IntOrStr):
+    reveal_type(int_or_str)  # revealed: Unknown
+```
+
+We *do* support stringified annotations if they appear in a position where a type expression is
+syntactically expected:
+
+```py
+ListOfInts = list["int"]
+
+def _(list_of_ints: ListOfInts):
+    reveal_type(list_of_ints)  # revealed: list[int]
+```
+
 ## Recursive
 
 ### Old union syntax

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

@@ -146,13 +146,11 @@ def _(flag: bool):
 def _(flag: bool):
     x = 1 if flag else "a"
 
-    # TODO: this should cause us to emit a diagnostic during
-    # type checking
+    # error: [invalid-argument-type] "Argument to function `isinstance` is incorrect: Expected `type | UnionType | tuple[Unknown, ...]`, found `Literal["a"]"
     if isinstance(x, "a"):
         reveal_type(x)  # revealed: Literal[1, "a"]
 
-    # TODO: this should cause us to emit a diagnostic during
-    # type checking
+    # error: [invalid-argument-type] "Argument to function `isinstance` is incorrect: Expected `type | UnionType | tuple[Unknown, ...]`, found `Literal["int"]"
     if isinstance(x, "int"):
         reveal_type(x)  # revealed: Literal[1, "a"]
 ```

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

@@ -214,8 +214,7 @@ def flag() -> bool:
 
 t = int if flag() else str
 
-# TODO: this should cause us to emit a diagnostic during
-# type checking
+# error: [invalid-argument-type] "Argument to function `issubclass` is incorrect: Expected `type | UnionType | tuple[Unknown, ...]`, found `Literal["str"]"
 if issubclass(t, "str"):
     reveal_type(t)  # revealed: <class 'int'> | <class 'str'>
 

crates/ty_python_semantic/src/types.rs

@@ -817,13 +817,11 @@ impl<'db> Type<'db> {
     }
 
     fn is_none(&self, db: &'db dyn Db) -> bool {
-        self.as_nominal_instance()
-            .is_some_and(|instance| instance.has_known_class(db, KnownClass::NoneType))
+        self.is_instance_of(db, KnownClass::NoneType)
     }
 
     fn is_bool(&self, db: &'db dyn Db) -> bool {
-        self.as_nominal_instance()
-            .is_some_and(|instance| instance.has_known_class(db, KnownClass::Bool))
+        self.is_instance_of(db, KnownClass::Bool)
     }
 
     fn is_enum(&self, db: &'db dyn Db) -> bool {
@@ -857,8 +855,7 @@ impl<'db> Type<'db> {
     }
 
     pub(crate) fn is_notimplemented(&self, db: &'db dyn Db) -> bool {
-        self.as_nominal_instance()
-            .is_some_and(|instance| instance.has_known_class(db, KnownClass::NotImplementedType))
+        self.is_instance_of(db, KnownClass::NotImplementedType)
     }
 
     pub(crate) const fn is_todo(&self) -> bool {
@@ -6404,6 +6401,17 @@ impl<'db> Type<'db> {
                     invalid_expressions: smallvec::smallvec_inline![InvalidTypeExpression::Generic],
                     fallback_type: Type::unknown(),
                 }),
+                KnownInstanceType::UnionType(union_type) => {
+                    let mut builder = UnionBuilder::new(db);
+                    for element in union_type.elements(db) {
+                        builder = builder.add(element.in_type_expression(
+                            db,
+                            scope_id,
+                            typevar_binding_context,
+                        )?);
+                    }
+                    Ok(builder.build())
+                }
             },
 
             Type::SpecialForm(special_form) => match special_form {
@@ -6572,9 +6580,6 @@ impl<'db> Type<'db> {
                 Some(KnownClass::GenericAlias) => Ok(todo_type!(
                     "Support for `typing.GenericAlias` instances in type expressions"
                 )),
-                Some(KnownClass::UnionType) => Ok(todo_type!(
-                    "Support for `types.UnionType` instances in type expressions"
-                )),
                 _ => Err(InvalidTypeExpressionError {
                     invalid_expressions: smallvec::smallvec_inline![
                         InvalidTypeExpression::InvalidType(*self, scope_id)
@@ -7614,6 +7619,10 @@ pub enum KnownInstanceType<'db> {
     /// A constraint set, which is exposed in mdtests as an instance of
     /// `ty_extensions.ConstraintSet`.
     ConstraintSet(TrackedConstraintSet<'db>),
+
+    /// A single instance of `types.UnionType`, which stores the left- and
+    /// right-hand sides of a PEP 604 union.
+    UnionType(UnionTypeInstance<'db>),
 }
 
 fn walk_known_instance_type<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
@@ -7640,6 +7649,11 @@ fn walk_known_instance_type<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
                 visitor.visit_type(db, default_ty);
             }
         }
+        KnownInstanceType::UnionType(union_type) => {
+            for element in union_type.elements(db) {
+                visitor.visit_type(db, *element);
+            }
+        }
     }
 }
 
@@ -7676,6 +7690,7 @@ impl<'db> KnownInstanceType<'db> {
                 // Nothing to normalize
                 Self::ConstraintSet(set)
             }
+            Self::UnionType(union_type) => Self::UnionType(union_type.normalized_impl(db, visitor)),
         }
     }
 
@@ -7690,6 +7705,7 @@ impl<'db> KnownInstanceType<'db> {
             Self::Deprecated(_) => KnownClass::Deprecated,
             Self::Field(_) => KnownClass::Field,
             Self::ConstraintSet(_) => KnownClass::ConstraintSet,
+            Self::UnionType(_) => KnownClass::UnionType,
         }
     }
 
@@ -7763,6 +7779,7 @@ impl<'db> KnownInstanceType<'db> {
                             constraints.display(self.db)
                         )
                     }
+                    KnownInstanceType::UnionType(_) => f.write_str("UnionType"),
                 }
             }
         }
@@ -8882,6 +8899,39 @@ impl<'db> TypeVarBoundOrConstraints<'db> {
     }
 }
 
+/// An instance of `types.UnionType`.
+///
+/// # Ordering
+/// Ordering is based on the context's salsa-assigned id and not on its values.
+/// The id may change between runs, or when the context was garbage collected and recreated.
+#[salsa::interned(debug)]
+#[derive(PartialOrd, Ord)]
+pub struct UnionTypeInstance<'db> {
+    #[returns(deref)]
+    elements: Box<[Type<'db>; 2]>,
+}
+
+impl get_size2::GetSize for UnionTypeInstance<'_> {}
+
+impl<'db> UnionTypeInstance<'db> {
+    pub(crate) fn from_elements(
+        db: &'db dyn Db,
+        left: Type<'db>,
+        right: Type<'db>,
+    ) -> UnionTypeInstance<'db> {
+        UnionTypeInstance::new(db, Box::new([left, right]))
+    }
+
+    pub(crate) fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
+        let elements = self.elements(db);
+        UnionTypeInstance::from_elements(
+            db,
+            elements[0].normalized_impl(db, visitor),
+            elements[1].normalized_impl(db, visitor),
+        )
+    }
+}
+
 /// Error returned if a type is not awaitable.
 #[derive(Debug)]
 enum AwaitError<'db> {

crates/ty_python_semantic/src/types/class.rs

@@ -1307,9 +1307,7 @@ impl<'db> Field<'db> {
     /// Returns true if this field is a `dataclasses.KW_ONLY` sentinel.
     /// <https://docs.python.org/3/library/dataclasses.html#dataclasses.KW_ONLY>
     pub(crate) fn is_kw_only_sentinel(&self, db: &'db dyn Db) -> bool {
-        self.declared_ty
-            .as_nominal_instance()
-            .is_some_and(|instance| instance.has_known_class(db, KnownClass::KwOnly))
+        self.declared_ty.is_instance_of(db, KnownClass::KwOnly)
     }
 }
 

crates/ty_python_semantic/src/types/class_base.rs

@@ -170,7 +170,8 @@ impl<'db> ClassBase<'db> {
                 | KnownInstanceType::TypeVar(_)
                 | KnownInstanceType::Deprecated(_)
                 | KnownInstanceType::Field(_)
-                | KnownInstanceType::ConstraintSet(_) => None,
+                | KnownInstanceType::ConstraintSet(_)
+                | KnownInstanceType::UnionType(_) => None,
             },
 
             Type::SpecialForm(special_form) => match special_form {