From 781b653511c920b3ecbdac851a0aa03c3a5f97b1 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 9 Apr 2025 18:33:16 +0100
Subject: [PATCH] [red-knot] Fix false positives on `types.UnionType` instances
 in type expressions (#17297)

---
 .../resources/mdtest/annotations/union.md     | 19 ++++++++++++++++++-
 crates/red_knot_python_semantic/src/types.rs  |  3 +++
 .../src/types/class.rs                        |  9 +++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/crates/red_knot_python_semantic/resources/mdtest/annotations/union.md b/crates/red_knot_python_semantic/resources/mdtest/annotations/union.md
index 253fe6dde6..e4d47a3a94 100644
--- a/crates/red_knot_python_semantic/resources/mdtest/annotations/union.md
+++ b/crates/red_knot_python_semantic/resources/mdtest/annotations/union.md
@@ -2,7 +2,7 @@
 
 ## Annotation
 
-`typing.Union` can be used to construct union types same as `|` operator.
+`typing.Union` can be used to construct union types in the same way as the `|` operator.
 
 ```py
 from typing import Union
@@ -69,3 +69,20 @@ from typing import Union
 def f(x: Union) -> None:
     reveal_type(x)  # revealed: Unknown
 ```
+
+## Implicit type aliases using new-style unions
+
+We don't recognise 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"
+```
+
+```py
+X = int | str
+
+def f(y: X):
+    reveal_type(y)  # revealed: @Todo(Support for `types.UnionType` instances in type expressions)
+```
diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index 3197f7fa9a..41657e3bb8 100644
--- a/crates/red_knot_python_semantic/src/types.rs
+++ b/crates/red_knot_python_semantic/src/types.rs
@@ -4035,6 +4035,9 @@ 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![InvalidTypeExpression::InvalidType(
                         *self
diff --git a/crates/red_knot_python_semantic/src/types/class.rs b/crates/red_knot_python_semantic/src/types/class.rs
index a3fe1024b4..aca5acd46d 100644
--- a/crates/red_knot_python_semantic/src/types/class.rs
+++ b/crates/red_knot_python_semantic/src/types/class.rs
@@ -1170,6 +1170,7 @@ pub(crate) enum KnownClass {
     MethodType,
     MethodWrapperType,
     WrapperDescriptorType,
+    UnionType,
     // Typeshed
     NoneType, // Part of `types` for Python >= 3.10
     // Typing
@@ -1233,6 +1234,7 @@ impl<'db> KnownClass {
             | Self::ParamSpecKwargs
             | Self::TypeVarTuple
             | Self::WrapperDescriptorType
+            | Self::UnionType
             | Self::MethodWrapperType => Truthiness::AlwaysTrue,
 
             Self::NoneType => Truthiness::AlwaysFalse,
@@ -1305,6 +1307,7 @@ impl<'db> KnownClass {
             Self::ModuleType => "ModuleType",
             Self::FunctionType => "FunctionType",
             Self::MethodType => "MethodType",
+            Self::UnionType => "UnionType",
             Self::MethodWrapperType => "MethodWrapperType",
             Self::WrapperDescriptorType => "WrapperDescriptorType",
             Self::NoneType => "NoneType",
@@ -1487,6 +1490,7 @@ impl<'db> KnownClass {
             | Self::FunctionType
             | Self::MethodType
             | Self::MethodWrapperType
+            | Self::UnionType
             | Self::WrapperDescriptorType => KnownModule::Types,
             Self::NoneType => KnownModule::Typeshed,
             Self::Any
@@ -1539,6 +1543,7 @@ impl<'db> KnownClass {
             | Self::VersionInfo
             | Self::EllipsisType
             | Self::TypeAliasType
+            | Self::UnionType
             | Self::NotImplementedType => true,
 
             Self::Any
@@ -1643,6 +1648,7 @@ impl<'db> KnownClass {
             | Self::Sized
             | Self::Enum
             | Self::Super
+            | Self::UnionType
             | Self::NewType => false,
         }
     }
@@ -1681,6 +1687,7 @@ impl<'db> KnownClass {
             "ModuleType" => Self::ModuleType,
             "FunctionType" => Self::FunctionType,
             "MethodType" => Self::MethodType,
+            "UnionType" => Self::UnionType,
             "MethodWrapperType" => Self::MethodWrapperType,
             "WrapperDescriptorType" => Self::WrapperDescriptorType,
             "NewType" => Self::NewType,
@@ -1758,6 +1765,7 @@ impl<'db> KnownClass {
             | Self::Enum
             | Self::Super
             | Self::NotImplementedType
+            | Self::UnionType
             | Self::WrapperDescriptorType => module == self.canonical_module(db),
             Self::NoneType => matches!(module, KnownModule::Typeshed | KnownModule::Types),
             Self::SpecialForm
@@ -2259,6 +2267,7 @@ mod tests {
 
         for class in KnownClass::iter() {
             let version_added = match class {
+                KnownClass::UnionType => PythonVersion::PY310,
                 KnownClass::BaseExceptionGroup => PythonVersion::PY311,
                 KnownClass::GenericAlias => PythonVersion::PY39,
                 _ => PythonVersion::PY37,