[ty] Allow PEP-604 unions in stubs and `TYPE_CHECKING` blocks prior to 3.10 (#21379)

crates/ty_python_semantic/resources/mdtest/implicit_type_aliases.md

@@ -272,6 +272,54 @@ def g(
 ): ...
 ```
 
+## `|` unions in stubs and `TYPE_CHECKING` blocks
+
+In runtime contexts, `|` unions are only permitted on Python 3.10+. But in suites of code that are
+never executed at runtime (stub files, `if TYPE_CHECKING` blocks, and stringified annotations), they
+are permitted even if the target version is set to Python 3.9 or earlier.
+
+```toml
+[environment]
+python-version = "3.9"
+```
+
+`bar.pyi`:
+
+```pyi
+Z = int | str
+GLOBAL_CONSTANT: Z
+```
+
+`foo.py`:
+
+```py
+from typing import TYPE_CHECKING
+from bar import GLOBAL_CONSTANT
+
+reveal_type(GLOBAL_CONSTANT)  # revealed: int | str
+
+if TYPE_CHECKING:
+    class ItsQuiteCloudyInManchester:
+        X = int | str
+
+        def f(obj: X):
+            reveal_type(obj)  # revealed: int | str
+
+    # TODO: we currently only understand code as being inside a `TYPE_CHECKING` block
+    # if a whole *scope* is inside the `if TYPE_CHECKING` block
+    # (like the `ItsQuiteCloudyInManchester` class above); this is a false-positive
+    Y = int | str  # error: [unsupported-operator]
+
+    def g(obj: Y):
+        # TODO: should be `int | str`
+        reveal_type(obj)  # revealed: Unknown
+
+Y = list["int | str"]
+
+def g(obj: Y):
+    reveal_type(obj)  # revealed: list[int | str]
+```
+
 ## Generic types
 
 Implicit type aliases can also refer to generic types:

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

@@ -8899,6 +8899,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             emitted_division_by_zero_diagnostic = self.check_division_by_zero(node, op, left_ty);
         }
 
+        let pep_604_unions_allowed = || {
+            Program::get(self.db()).python_version(self.db()) >= PythonVersion::PY310
+                || self.file().is_stub(self.db())
+                || self.scope().scope(self.db()).in_type_checking_block()
+        };
+
         match (left_ty, right_ty, op) {
             (Type::Union(lhs_union), rhs, _) => lhs_union.try_map(self.db(), |lhs_element| {
                 self.infer_binary_expression_type(
@@ -9160,7 +9166,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     | KnownInstanceType::Annotated(_),
                 ),
                 ast::Operator::BitOr,
-            ) if Program::get(self.db()).python_version(self.db()) >= PythonVersion::PY310 => {
+            ) if pep_604_unions_allowed() => {
                 if left_ty.is_equivalent_to(self.db(), right_ty) {
                     Some(left_ty)
                 } else {
@@ -9186,7 +9192,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 | Type::KnownInstance(..)
                 | Type::SpecialForm(..),
                 ast::Operator::BitOr,
-            ) if Program::get(self.db()).python_version(self.db()) >= PythonVersion::PY310
+            ) if pep_604_unions_allowed()
                 && instance.has_known_class(self.db(), KnownClass::NoneType) =>
             {
                 Some(Type::KnownInstance(KnownInstanceType::UnionType(
@@ -9210,8 +9216,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 _,
                 Type::ClassLiteral(..) | Type::GenericAlias(..) | Type::SubclassOf(..),
                 ast::Operator::BitOr,
-            ) if Program::get(self.db()).python_version(self.db()) >= PythonVersion::PY310 => {
-                Type::try_call_bin_op_with_policy(
+            ) if pep_604_unions_allowed() => Type::try_call_bin_op_with_policy(
                 self.db(),
                 left_ty,
                 ast::Operator::BitOr,
@@ -9219,8 +9224,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 MemberLookupPolicy::META_CLASS_NO_TYPE_FALLBACK,
             )
             .ok()
-                .map(|binding| binding.return_type(self.db()))
-            }
+            .map(|binding| binding.return_type(self.db())),
 
             // We've handled all of the special cases that we support for literals, so we need to
             // fall back on looking for dunder methods on one of the operand types.