[ty] Add narrowing for isinstance() and issubclass() checks that use PEP-604 unions (#21334)

Author: AlexWaygood

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

@@ -70,6 +70,74 @@ def _(flag: bool):
         reveal_type(x)  # revealed: Literal["a"]
 ```
 
+## `classinfo` is a PEP-604 union of types
+
+```toml
+[environment]
+python-version = "3.10"
+```
+
+```py
+def _(x: int | str | bytes | memoryview | range):
+    if isinstance(x, int | str):
+        reveal_type(x)  # revealed: int | str
+    elif isinstance(x, bytes | memoryview):
+        reveal_type(x)  # revealed: bytes | memoryview[Unknown]
+    else:
+        reveal_type(x)  # revealed: range
+```
+
+Although `isinstance()` usually only works if all elements in the `UnionType` are class objects, at
+runtime a special exception is made for `None` so that `isinstance(x, int | None)` can work:
+
+```py
+def _(x: int | str | bytes | range | None):
+    if isinstance(x, int | str | None):
+        reveal_type(x)  # revealed: int | str | None
+    else:
+        reveal_type(x)  # revealed: bytes | range
+```
+
+## `classinfo` is an invalid PEP-604 union of types
+
+Except for the `None` special case mentioned above, narrowing can only take place if all elements in
+the PEP-604 union are class literals. If any elements are generic aliases or other types, the
+`isinstance()` call may fail at runtime, so no narrowing can take place:
+
+```toml
+[environment]
+python-version = "3.10"
+```
+
+```py
+def _(x: int | list[int] | bytes):
+    # TODO: this fails at runtime; we should emit a diagnostic
+    # (requires special-casing of the `isinstance()` signature)
+    if isinstance(x, int | list[int]):
+        reveal_type(x)  # revealed: int | list[int] | bytes
+    else:
+        reveal_type(x)  # revealed: int | list[int] | bytes
+```
+
+## PEP-604 unions on Python \<3.10
+
+PEP-604 unions were added in Python 3.10, so attempting to use them on Python 3.9 does not lead to
+any type narrowing.
+
+```toml
+[environment]
+python-version = "3.9"
+```
+
+```py
+def _(x: int | str | bytes):
+    # error: [unsupported-operator]
+    if isinstance(x, int | str):
+        reveal_type(x)  # revealed: (int & Unknown) | (str & Unknown) | (bytes & Unknown)
+    else:
+        reveal_type(x)  # revealed: (int & Unknown) | (str & Unknown) | (bytes & Unknown)
+```
+
 ## Class types
 
 ```py

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

@@ -131,6 +131,74 @@ def _(flag1: bool, flag2: bool):
         reveal_type(t)  # revealed: <class 'str'>
 ```
 
+## `classinfo` is a PEP-604 union of types
+
+```toml
+[environment]
+python-version = "3.10"
+```
+
+```py
+def f(x: type[int | str | bytes | range]):
+    if issubclass(x, int | str):
+        reveal_type(x)  # revealed: type[int] | type[str]
+    elif issubclass(x, bytes | memoryview):
+        reveal_type(x)  # revealed: type[bytes]
+    else:
+        reveal_type(x)  # revealed: <class 'range'>
+```
+
+Although `issubclass()` usually only works if all elements in the `UnionType` are class objects, at
+runtime a special exception is made for `None` so that `issubclass(x, int | None)` can work:
+
+```py
+def _(x: type):
+    if issubclass(x, int | str | None):
+        reveal_type(x)  # revealed: type[int] | type[str] | <class 'NoneType'>
+    else:
+        reveal_type(x)  # revealed: type & ~type[int] & ~type[str] & ~<class 'NoneType'>
+```
+
+## `classinfo` is an invalid PEP-604 union of types
+
+Except for the `None` special case mentioned above, narrowing can only take place if all elements in
+the PEP-604 union are class literals. If any elements are generic aliases or other types, the
+`issubclass()` call may fail at runtime, so no narrowing can take place:
+
+```toml
+[environment]
+python-version = "3.10"
+```
+
+```py
+def _(x: type[int | list | bytes]):
+    # TODO: this fails at runtime; we should emit a diagnostic
+    # (requires special-casing of the `issubclass()` signature)
+    if issubclass(x, int | list[int]):
+        reveal_type(x)  # revealed: type[int] | type[list[Unknown]] | type[bytes]
+    else:
+        reveal_type(x)  # revealed: type[int] | type[list[Unknown]] | type[bytes]
+```
+
+## PEP-604 unions on Python \<3.10
+
+PEP-604 unions were added in Python 3.10, so attempting to use them on Python 3.9 does not lead to
+any type narrowing.
+
+```toml
+[environment]
+python-version = "3.9"
+```
+
+```py
+def _(x: type[int | str | bytes]):
+    # error: [unsupported-operator]
+    if issubclass(x, int | str):
+        reveal_type(x)  # revealed: (type[int] & Unknown) | (type[str] & Unknown) | (type[bytes] & Unknown)
+    else:
+        reveal_type(x)  # revealed: (type[int] & Unknown) | (type[str] & Unknown) | (type[bytes] & Unknown)
+```
+
 ## Special cases
 
 ### Emit a diagnostic if the first argument is of wrong type

crates/ty_python_semantic/src/types/narrow.rs

@@ -11,9 +11,9 @@ use crate::types::enums::{enum_member_literals, enum_metadata};
 use crate::types::function::KnownFunction;
 use crate::types::infer::infer_same_file_expression_type;
 use crate::types::{
-    ClassLiteral, ClassType, IntersectionBuilder, KnownClass, SpecialFormType, SubclassOfInner,
-    SubclassOfType, Truthiness, Type, TypeContext, TypeVarBoundOrConstraints, UnionBuilder,
-    infer_expression_types,
+    ClassLiteral, ClassType, IntersectionBuilder, KnownClass, KnownInstanceType, SpecialFormType,
+    SubclassOfInner, SubclassOfType, Truthiness, Type, TypeContext, TypeVarBoundOrConstraints,
+    UnionBuilder, infer_expression_types,
 };
 
 use ruff_db::parsed::{ParsedModuleRef, parsed_module};
@@ -212,6 +212,23 @@ impl ClassInfoConstraintFunction {
                 )
             }),
 
+            Type::KnownInstance(KnownInstanceType::UnionType(elements)) => {
+                UnionType::try_from_elements(
+                    db,
+                    elements.elements(db).iter().map(|element| {
+                        // A special case is made for `None` at runtime
+                        // (it's implicitly converted to `NoneType` in `int | None`)
+                        // which means that `isinstance(x, int | None)` works even though
+                        // `None` is not a class literal.
+                        if element.is_none(db) {
+                            self.generate_constraint(db, KnownClass::NoneType.to_class_literal(db))
+                        } else {
+                            self.generate_constraint(db, *element)
+                        }
+                    }),
+                )
+            }
+
             Type::AlwaysFalsy
             | Type::AlwaysTruthy
             | Type::BooleanLiteral(_)