[ty] implement typing.TypeGuard (#20974)

## Summary

Resolve(s) astral-sh/ty#117, astral-sh/ty#1569

Implement `typing.TypeGuard`. Due to the fact that it [overrides
anything previously known about the checked
value](https://typing.python.org/en/latest/spec/narrowing.html#typeguard)---

> When a conditional statement includes a call to a user-defined type
guard function, and that function returns true, the expression passed as
the first positional argument to the type guard function should be
assumed by a static type checker to take on the type specified in the
TypeGuard return type, unless and until it is further narrowed within
the conditional code block.

---we have to substantially rework the constraints system. In
particular, we make constraints represented as a disjunctive normal form
(DNF) where each term includes a regular constraint, and one or more
disjuncts with a typeguard constraint. Some test cases (including some
with more complex boolean logic) are added to `type_guards.md`.


## Test Plan

- update existing tests
- add new tests for more complex boolean logic with `TypeGuard`
- add new tests for `TypeGuard` variance

---------

Co-authored-by: Carl Meyer <carl@astral.sh>

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

@@ -16,7 +16,6 @@ def f(*args: Unpack[Ts]) -> tuple[Unpack[Ts]]:
     reveal_type(args)  # revealed: tuple[@Todo(`Unpack[]` special form), ...]
     return args
 
-def g() -> TypeGuard[int]: ...
 def i(callback: Callable[Concatenate[int, P], R_co], *args: P.args, **kwargs: P.kwargs) -> R_co:
     reveal_type(args)  # revealed: P@i.args
     reveal_type(kwargs)  # revealed: P@i.kwargs

crates/ty_python_semantic/resources/mdtest/generics/pep695/variance.md

@@ -790,6 +790,44 @@ static_assert(not is_assignable_to(C[B], C[A]))
 static_assert(not is_assignable_to(C[A], C[B]))
 ```
 
+## TypeGuard
+
+`TypeGuard[T]` is covariant in `T`. The typing spec doesn't explicitly call this out, but it follows
+from similar logic to invariance of `TypeIs` except without the negative case.
+
+Formally, suppose we have types `A` and `B` with `B < A`. Take `x: object` to be the value that all
+subsequent `TypeGuard`s are narrowing.
+
+We can assign `p: TypeGuard[A] = q` where `q: TypeGuard[B]` because
+
+- if `q` is `False`, then no constraints were learned on `x` before and none are now learned, so
+    nothing changes
+- if `q` is `True`, then we know `x: B`. From `B < A`, we conclude `x: A`.
+
+We _cannot_ assign `p: TypeGuard[B] = q` where `q: TypeGuard[A]` because if `q` is `True`, we would
+be concluding `x: B` from `x: A`, which is an unsafe downcast.
+
+```py
+from typing import TypeGuard
+from ty_extensions import is_assignable_to, is_subtype_of, static_assert
+
+class A:
+    pass
+
+class B(A):
+    pass
+
+class C[T]:
+    def check(x: object) -> TypeGuard[T]:
+        # this is a bad check, but we only care about it type-checking
+        return False
+
+static_assert(is_subtype_of(C[B], C[A]))
+static_assert(not is_subtype_of(C[A], C[B]))
+static_assert(is_assignable_to(C[B], C[A]))
+static_assert(not is_assignable_to(C[A], C[B]))
+```
+
 ## Type aliases
 
 The variance of the type alias matches the variance of the value type (RHS type).

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

@@ -12,21 +12,19 @@ from typing_extensions import TypeGuard, TypeIs
 def _(
     a: TypeGuard[str],
     b: TypeIs[str | int],
-    c: TypeGuard[Intersection[complex, Not[int], Not[float]]],
+    c: TypeGuard[bool],
     d: TypeIs[tuple[TypeOf[bytes]]],
     e: TypeGuard,  # error: [invalid-type-form]
     f: TypeIs,  # error: [invalid-type-form]
 ):
-    # TODO: Should be `TypeGuard[str]`
-    reveal_type(a)  # revealed: @Todo(`TypeGuard[]` special form)
+    reveal_type(a)  # revealed: TypeGuard[str]
     reveal_type(b)  # revealed: TypeIs[str | int]
-    # TODO: Should be `TypeGuard[complex & ~int & ~float]`
-    reveal_type(c)  # revealed: @Todo(`TypeGuard[]` special form)
+    reveal_type(c)  # revealed: TypeGuard[bool]
     reveal_type(d)  # revealed: TypeIs[tuple[<class 'bytes'>]]
     reveal_type(e)  # revealed: Unknown
     reveal_type(f)  # revealed: Unknown
 
-# TODO: error: [invalid-return-type] "Function always implicitly returns `None`, which is not assignable to return type `TypeGuard[str]`"
+# error: [invalid-return-type] "Function always implicitly returns `None`, which is not assignable to return type `TypeGuard[str]`"
 def _(a) -> TypeGuard[str]: ...
 
 # error: [invalid-return-type] "Function always implicitly returns `None`, which is not assignable to return type `TypeIs[str]`"
@@ -38,8 +36,7 @@ def g(a) -> TypeIs[str]:
     return True
 
 def _(a: object):
-    # TODO: Should be `TypeGuard[str @ a]`
-    reveal_type(f(a))  # revealed: @Todo(`TypeGuard[]` special form)
+    reveal_type(f(a))  # revealed: TypeGuard[str @ a]
     reveal_type(g(a))  # revealed: TypeIs[str @ a]
 ```
 
@@ -96,6 +93,72 @@ def _(a: int) -> TypeIs[str]: ...
 def _(a: bool | str) -> TypeIs[int]: ...
 ```
 
+## Methods
+
+Methods narrow the first positional argument after `self` or `cls`
+
+```py
+from typing import TypeGuard
+
+class C:
+    def f(self, x: object) -> TypeGuard[str]:
+        return True
+
+    @classmethod
+    def g(cls, x: object) -> TypeGuard[int]:
+        return True
+    # TODO: this could error at definition time
+    def h(self) -> TypeGuard[str]:
+        return True
+    # TODO: this could error at definition time
+    @classmethod
+    def j(cls) -> TypeGuard[int]:
+        return True
+
+def _(x: object):
+    if C().f(x):
+        reveal_type(x)  # revealed: str
+    if C.f(C(), x):
+        # TODO: should be str
+        reveal_type(x)  # revealed: object
+    if C.g(x):
+        reveal_type(x)  # revealed: int
+    if C().g(x):
+        reveal_type(x)  # revealed: int
+    if C().h():  # error: [invalid-type-guard-call] "Type guard call does not have a target"
+        pass
+    if C.j():  # error: [invalid-type-guard-call] "Type guard call does not have a target"
+        pass
+```
+
+```py
+from typing_extensions import TypeIs
+
+def is_int(val: object) -> TypeIs[int]:
+    return isinstance(val, int)
+
+class A:
+    def is_int(self, val: object) -> TypeIs[int]:
+        return isinstance(val, int)
+
+    @classmethod
+    def is_int2(cls, val: object) -> TypeIs[int]:
+        return isinstance(val, int)
+
+def _(x: object):
+    if is_int(x):
+        reveal_type(x)  # revealed: int
+
+    if A().is_int(x):
+        reveal_type(x)  # revealed: int
+
+    if A().is_int2(x):
+        reveal_type(x)  # revealed: int
+
+    if A.is_int2(x):
+        reveal_type(x)  # revealed: int
+```
+
 ## Arguments to special forms
 
 `TypeGuard` and `TypeIs` accept exactly one type argument.
@@ -105,15 +168,14 @@ from typing_extensions import TypeGuard, TypeIs
 
 a = 123
 
-# TODO: error: [invalid-type-form]
+# error: [invalid-type-form] "Special form `typing.TypeGuard` expected exactly one type parameter"
 def f(_) -> TypeGuard[int, str]: ...
 
 # error: [invalid-type-form] "Special form `typing.TypeIs` expected exactly one type parameter"
 # error: [invalid-type-form] "Variable of type `Literal[123]` is not allowed in a type expression"
 def g(_) -> TypeIs[a, str]: ...
 
-# TODO: Should be `Unknown`
-reveal_type(f(0))  # revealed: @Todo(`TypeGuard[]` special form)
+reveal_type(f(0))  # revealed: Unknown
 reveal_type(g(0))  # revealed: Unknown
 ```
 
@@ -126,9 +188,10 @@ from typing_extensions import Literal, TypeGuard, TypeIs, assert_never
 
 def _(a: object, flag: bool) -> TypeGuard[str]:
     if flag:
+        # error: [invalid-return-type] "Return type does not match returned value: expected `TypeGuard[str]`, found `Literal[0]`"
         return 0
 
-    # TODO: error: [invalid-return-type] "Return type does not match returned value: expected `TypeIs[str]`, found `Literal["foo"]`"
+    # error: [invalid-return-type] "Return type does not match returned value: expected `TypeGuard[str]`, found `Literal["foo"]`"
     return "foo"
 
 # error: [invalid-return-type] "Function can implicitly return `None`, which is not assignable to return type `TypeIs[str]`"
@@ -193,8 +256,7 @@ def is_bar(a: object) -> TypeIs[Bar]:
 
 def _(a: Foo | Bar):
     if guard_foo(a):
-        # TODO: Should be `Foo`
-        reveal_type(a)  # revealed: Foo | Bar
+        reveal_type(a)  # revealed: Foo
     else:
         reveal_type(a)  # revealed: Foo | Bar
 
@@ -204,6 +266,26 @@ def _(a: Foo | Bar):
         reveal_type(a)  # revealed: Foo & ~Bar
 ```
 
+```py
+from typing import TypeGuard, reveal_type
+
+class P:
+    pass
+
+class A:
+    pass
+
+class B:
+    pass
+
+def is_b(val: object) -> TypeGuard[B]:
+    return isinstance(val, B)
+
+def _(x: P):
+    if isinstance(x, A) or is_b(x):
+        reveal_type(x)  # revealed: B | (P & A)
+```
+
 Attribute and subscript narrowing is supported:
 
 ```py
@@ -215,23 +297,17 @@ class C(Generic[T]):
     v: T
 
 def _(a: tuple[Foo, Bar] | tuple[Bar, Foo], c: C[Any]):
-    # TODO: Should be `TypeGuard[Foo @ a[1]]`
-    if reveal_type(guard_foo(a[1])):  # revealed: @Todo(`TypeGuard[]` special form)
-        # TODO: Should be `tuple[Bar, Foo]`
+    if reveal_type(guard_foo(a[1])):  # revealed: TypeGuard[Foo @ a[1]]
         reveal_type(a)  # revealed: tuple[Foo, Bar] | tuple[Bar, Foo]
-        # TODO: Should be `Foo`
-        reveal_type(a[1])  # revealed: Bar | Foo
+        reveal_type(a[1])  # revealed: Foo
 
     if reveal_type(is_bar(a[0])):  # revealed: TypeIs[Bar @ a[0]]
-        # TODO: Should be `tuple[Bar, Bar & Foo]`
         reveal_type(a)  # revealed: tuple[Foo, Bar] | tuple[Bar, Foo]
         reveal_type(a[0])  # revealed: Bar
 
-    # TODO: Should be `TypeGuard[Foo @ c.v]`
-    if reveal_type(guard_foo(c.v)):  # revealed: @Todo(`TypeGuard[]` special form)
+    if reveal_type(guard_foo(c.v)):  # revealed: TypeGuard[Foo @ c.v]
         reveal_type(c)  # revealed: C[Any]
-        # TODO: Should be `Foo`
-        reveal_type(c.v)  # revealed: Any
+        reveal_type(c.v)  # revealed: Foo
 
     if reveal_type(is_bar(c.v)):  # revealed: TypeIs[Bar @ c.v]
         reveal_type(c)  # revealed: C[Any]
@@ -246,8 +322,7 @@ def _(a: Foo | Bar):
     c = is_bar(a)
 
     reveal_type(a)  # revealed: Foo | Bar
-    # TODO: Should be `TypeGuard[Foo @ a]`
-    reveal_type(b)  # revealed: @Todo(`TypeGuard[]` special form)
+    reveal_type(b)  # revealed: TypeGuard[Foo @ a]
     reveal_type(c)  # revealed: TypeIs[Bar @ a]
 
     if b:
@@ -345,25 +420,82 @@ class Baz(Bar): ...
 def guard_foo(a: object) -> TypeGuard[Foo]:
     return True
 
+def guard_bar(a: object) -> TypeGuard[Bar]:
+    return True
+
 def is_bar(a: object) -> TypeIs[Bar]:
     return True
 
 def does_not_narrow_in_negative_case(a: Foo | Bar):
     if not guard_foo(a):
-        # TODO: Should be `Bar`
         reveal_type(a)  # revealed: Foo | Bar
     else:
-        reveal_type(a)  # revealed: Foo | Bar
+        reveal_type(a)  # revealed: Foo
 
 def narrowed_type_must_be_exact(a: object, b: Baz):
     if guard_foo(b):
-        # TODO: Should be `Foo`
-        reveal_type(b)  # revealed: Baz
+        reveal_type(b)  # revealed: Foo
 
     if isinstance(a, Baz) and is_bar(a):
         reveal_type(a)  # revealed: Baz
 
     if isinstance(a, Bar) and guard_foo(a):
-        # TODO: Should be `Foo`
-        reveal_type(a)  # revealed: Bar
+        reveal_type(a)  # revealed: Foo
+        if guard_bar(a):
+            reveal_type(a)  # revealed: Bar
+```
+
+## TypeGuard overrides normal constraints
+
+TypeGuard constraints override any previous narrowing, but additional "regular" constraints can be
+added on to TypeGuard constraints.
+
+```py
+from typing_extensions import TypeGuard, TypeIs
+
+class A: ...
+class B: ...
+class C: ...
+
+def f(x: object) -> TypeGuard[A]:
+    return True
+
+def g(x: object) -> TypeGuard[B]:
+    return True
+
+def h(x: object) -> TypeIs[C]:
+    return True
+
+def _(x: object):
+    if f(x) and g(x) and h(x):
+        reveal_type(x)  # revealed: B & C
+```
+
+## Boolean logic with TypeGuard and TypeIs
+
+TypeGuard constraints need to properly distribute through boolean operations.
+
+```py
+from typing_extensions import TypeGuard, TypeIs
+
+class A: ...
+class B: ...
+class C: ...
+
+def f(x: object) -> TypeIs[A]:
+    return True
+
+def g(x: object) -> TypeGuard[B]:
+    return True
+
+def h(x: object) -> TypeIs[C]:
+    return True
+
+def _(x: object):
+    # g(x) or h(x) should give B | C
+    # Then f(x) and (...) should distribute: (f(x) and g(x)) or (f(x) and h(x))
+    # Which is (Regular(A) & TypeGuard(B)) | (Regular(A) & Regular(C))
+    # TypeGuard clobbers in the first branch, giving: B | (A & C)
+    if f(x) and (g(x) or h(x)):
+        reveal_type(x)  # revealed: B | (A & C)
 ```

crates/ty_python_semantic/resources/mdtest/type_properties/is_assignable_to.md

@@ -1383,8 +1383,7 @@ from typing_extensions import Any, TypeGuard, TypeIs
 static_assert(is_assignable_to(TypeGuard[Unknown], bool))
 static_assert(is_assignable_to(TypeIs[Any], bool))
 
-# TODO no error
-static_assert(not is_assignable_to(TypeGuard[Unknown], str))  # error: [static-assert-error]
+static_assert(not is_assignable_to(TypeGuard[Unknown], str))
 static_assert(not is_assignable_to(TypeIs[Any], str))
 ```
 

crates/ty_python_semantic/resources/mdtest/type_properties/is_disjoint_from.md

@@ -578,8 +578,7 @@ from typing_extensions import TypeGuard, TypeIs
 static_assert(not is_disjoint_from(bool, TypeGuard[str]))
 static_assert(not is_disjoint_from(bool, TypeIs[str]))
 
-# TODO no error
-static_assert(is_disjoint_from(str, TypeGuard[str]))  # error: [static-assert-error]
+static_assert(is_disjoint_from(str, TypeGuard[str]))
 static_assert(is_disjoint_from(str, TypeIs[str]))
 ```
 

crates/ty_python_semantic/resources/mdtest/type_properties/is_subtype_of.md

@@ -670,9 +670,8 @@ Fully-static `TypeGuard[...]` and `TypeIs[...]` are subtypes of `bool`.
 from ty_extensions import is_subtype_of, static_assert
 from typing_extensions import TypeGuard, TypeIs
 
-# TODO: TypeGuard
-# static_assert(is_subtype_of(TypeGuard[int], bool))
-# static_assert(is_subtype_of(TypeGuard[int], int))
+static_assert(is_subtype_of(TypeGuard[str], bool))
+static_assert(is_subtype_of(TypeGuard[str], int))
 static_assert(is_subtype_of(TypeIs[str], bool))
 static_assert(is_subtype_of(TypeIs[str], int))
 ```
@@ -683,12 +682,12 @@ static_assert(is_subtype_of(TypeIs[str], int))
 from ty_extensions import is_equivalent_to, is_subtype_of, static_assert
 from typing_extensions import TypeGuard, TypeIs
 
-# TODO: TypeGuard
-# static_assert(is_subtype_of(TypeGuard[int], TypeGuard[int]))
-# static_assert(is_subtype_of(TypeGuard[bool], TypeGuard[int]))
+static_assert(is_subtype_of(TypeGuard[int], TypeGuard[int]))
+static_assert(is_subtype_of(TypeGuard[bool], TypeGuard[int]))
 static_assert(is_subtype_of(TypeIs[int], TypeIs[int]))
 static_assert(is_subtype_of(TypeIs[int], TypeIs[int]))
 
+static_assert(is_subtype_of(TypeGuard[bool], TypeGuard[int]))
 static_assert(not is_subtype_of(TypeGuard[int], TypeGuard[bool]))
 static_assert(not is_subtype_of(TypeIs[bool], TypeIs[int]))
 static_assert(not is_subtype_of(TypeIs[int], TypeIs[bool]))