[ty] Implement constraint implication for compound types

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

@@ -12,8 +12,7 @@ a particular constraint set hold_.
 ## Concrete types
 
 For concrete types, constraint implication is exactly the same as subtyping. (A concrete type is any
-fully static type that is not a typevar. It can _contain_ a typevar, though — `list[T]` is
-considered concrete.)
+fully static type that does not contain a typevar.)
 
 ```py
 from ty_extensions import ConstraintSet, is_subtype_of, static_assert
@@ -191,4 +190,163 @@ def mutually_constrained[T, U]():
     static_assert(not given_int.implies_subtype_of(T, str))
 ```
 
+## Compound types
+
+All of the relationships in the above section also apply when a typevar appears in a compound type.
+
+```py
+from typing import Never
+from ty_extensions import ConstraintSet, static_assert
+
+class Covariant[T]:
+    def get(self) -> T:
+        raise ValueError
+
+def given_constraints[T]():
+    static_assert(not ConstraintSet.always().implies_subtype_of(Covariant[T], Covariant[int]))
+    static_assert(not ConstraintSet.always().implies_subtype_of(Covariant[T], Covariant[bool]))
+    static_assert(not ConstraintSet.always().implies_subtype_of(Covariant[T], Covariant[str]))
+
+    # These are vacuously true; false implies anything
+    static_assert(ConstraintSet.never().implies_subtype_of(Covariant[T], Covariant[int]))
+    static_assert(ConstraintSet.never().implies_subtype_of(Covariant[T], Covariant[bool]))
+    static_assert(ConstraintSet.never().implies_subtype_of(Covariant[T], Covariant[str]))
+
+    # For a covariant typevar, (T ≤ int) implies that (Covariant[T] ≤ Covariant[int]).
+    given_int = ConstraintSet.range(Never, T, int)
+    static_assert(given_int.implies_subtype_of(Covariant[T], Covariant[int]))
+    static_assert(not given_int.implies_subtype_of(Covariant[T], Covariant[bool]))
+    static_assert(not given_int.implies_subtype_of(Covariant[T], Covariant[str]))
+
+    given_bool = ConstraintSet.range(Never, T, bool)
+    static_assert(given_bool.implies_subtype_of(Covariant[T], Covariant[int]))
+    static_assert(given_bool.implies_subtype_of(Covariant[T], Covariant[bool]))
+    static_assert(not given_bool.implies_subtype_of(Covariant[T], Covariant[str]))
+
+    given_bool_int = ConstraintSet.range(bool, T, int)
+    static_assert(not given_bool_int.implies_subtype_of(Covariant[int], Covariant[T]))
+    static_assert(given_bool_int.implies_subtype_of(Covariant[bool], Covariant[T]))
+    static_assert(not given_bool_int.implies_subtype_of(Covariant[str], Covariant[T]))
+
+def mutually_constrained[T, U]():
+    # If (T = U ∧ U ≤ int), then (T ≤ int) must be true as well, and therefore
+    # (Covariant[T] ≤ Covariant[int]).
+    given_int = ConstraintSet.range(U, T, U) & ConstraintSet.range(Never, U, int)
+    static_assert(given_int.implies_subtype_of(Covariant[T], Covariant[int]))
+    static_assert(not given_int.implies_subtype_of(Covariant[T], Covariant[bool]))
+    static_assert(not given_int.implies_subtype_of(Covariant[T], Covariant[str]))
+
+    # If (T ≤ U ∧ U ≤ int), then (T ≤ int) must be true as well, and therefore
+    # (Covariant[T] ≤ Covariant[int]).
+    given_int = ConstraintSet.range(Never, T, U) & ConstraintSet.range(Never, U, int)
+    static_assert(given_int.implies_subtype_of(Covariant[T], Covariant[int]))
+    static_assert(not given_int.implies_subtype_of(Covariant[T], Covariant[bool]))
+    static_assert(not given_int.implies_subtype_of(Covariant[T], Covariant[str]))
+```
+
+Many of the relationships are reversed for typevars that appear in contravariant types.
+
+```py
+class Contravariant[T]:
+    def set(self, value: T):
+        pass
+
+def given_constraints[T]():
+    static_assert(not ConstraintSet.always().implies_subtype_of(Contravariant[int], Contravariant[T]))
+    static_assert(not ConstraintSet.always().implies_subtype_of(Contravariant[bool], Contravariant[T]))
+    static_assert(not ConstraintSet.always().implies_subtype_of(Contravariant[str], Contravariant[T]))
+
+    # These are vacuously true; false implies anything
+    static_assert(ConstraintSet.never().implies_subtype_of(Contravariant[int], Contravariant[T]))
+    static_assert(ConstraintSet.never().implies_subtype_of(Contravariant[bool], Contravariant[T]))
+    static_assert(ConstraintSet.never().implies_subtype_of(Contravariant[str], Contravariant[T]))
+
+    # For a contravariant typevar, (T ≤ int) implies that (Contravariant[int] ≤ Contravariant[T]).
+    # (The order of the comparison is reversed because of contravariance.)
+    given_int = ConstraintSet.range(Never, T, int)
+    static_assert(given_int.implies_subtype_of(Contravariant[int], Contravariant[T]))
+    static_assert(not given_int.implies_subtype_of(Contravariant[bool], Contravariant[T]))
+    static_assert(not given_int.implies_subtype_of(Contravariant[str], Contravariant[T]))
+
+    given_bool = ConstraintSet.range(Never, T, int)
+    static_assert(given_bool.implies_subtype_of(Contravariant[int], Contravariant[T]))
+    static_assert(not given_bool.implies_subtype_of(Contravariant[bool], Contravariant[T]))
+    static_assert(not given_bool.implies_subtype_of(Contravariant[str], Contravariant[T]))
+
+def mutually_constrained[T, U]():
+    # If (T = U ∧ U ≤ int), then (T ≤ int) must be true as well, and therefore
+    # (Contravariant[int] ≤ Contravariant[T]).
+    given_int = ConstraintSet.range(U, T, U) & ConstraintSet.range(Never, U, int)
+    static_assert(given_int.implies_subtype_of(Contravariant[int], Contravariant[T]))
+    static_assert(not given_int.implies_subtype_of(Contravariant[bool], Contravariant[T]))
+    static_assert(not given_int.implies_subtype_of(Contravariant[str], Contravariant[T]))
+
+    # If (T ≤ U ∧ U ≤ int), then (T ≤ int) must be true as well, and therefore
+    # (Contravariant[int] ≤ Contravariant[T]).
+    given_int = ConstraintSet.range(Never, T, U) & ConstraintSet.range(Never, U, int)
+    static_assert(given_int.implies_subtype_of(Contravariant[int], Contravariant[T]))
+    static_assert(not given_int.implies_subtype_of(Contravariant[bool], Contravariant[T]))
+    static_assert(not given_int.implies_subtype_of(Contravariant[str], Contravariant[T]))
+```
+
+For invariant typevars, subtyping of the typevar does not imply subtyping of the compound type in
+either direction. But an equality constraint on the typevar does.
+
+```py
+class Invariant[T]:
+    def get(self) -> T:
+        raise ValueError
+
+    def set(self, value: T):
+        pass
+
+def given_constraints[T]():
+    static_assert(not ConstraintSet.always().implies_subtype_of(Invariant[T], Invariant[int]))
+    static_assert(not ConstraintSet.always().implies_subtype_of(Invariant[T], Invariant[bool]))
+    static_assert(not ConstraintSet.always().implies_subtype_of(Invariant[T], Invariant[str]))
+
+    # These are vacuously true; false implies anything
+    static_assert(ConstraintSet.never().implies_subtype_of(Invariant[T], Invariant[int]))
+    static_assert(ConstraintSet.never().implies_subtype_of(Invariant[T], Invariant[bool]))
+    static_assert(ConstraintSet.never().implies_subtype_of(Invariant[T], Invariant[str]))
+
+    # For an invariant typevar, (T ≤ int) does not imply that (Invariant[T] ≤ Invariant[int]).
+    given_int = ConstraintSet.range(Never, T, int)
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[int]))
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[bool]))
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[str]))
+
+    # It also does not imply the contravariant ordering (Invariant[int] ≤ Invariant[T]).
+    static_assert(not given_int.implies_subtype_of(Invariant[int], Invariant[T]))
+    static_assert(not given_int.implies_subtype_of(Invariant[bool], Invariant[T]))
+    static_assert(not given_int.implies_subtype_of(Invariant[str], Invariant[T]))
+
+    # But (T = int) does imply both.
+    given_int = ConstraintSet.range(int, T, int)
+    static_assert(given_int.implies_subtype_of(Invariant[T], Invariant[int]))
+    static_assert(given_int.implies_subtype_of(Invariant[int], Invariant[T]))
+    static_assert(not given_int.implies_subtype_of(Invariant[bool], Invariant[T]))
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[bool]))
+    static_assert(not given_int.implies_subtype_of(Invariant[str], Invariant[T]))
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[str]))
+
+def mutually_constrained[T, U]():
+    # If (T = U ∧ U ≤ int), then (T ≤ int) must be true as well. But because T is invariant, that
+    # does _not_ imply that (Invariant[T] ≤ Invariant[int]).
+    given_int = ConstraintSet.range(U, T, U) & ConstraintSet.range(Never, U, int)
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[int]))
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[bool]))
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[str]))
+
+    # If (T = U ∧ U = int), then (T = int) must be true as well. That is an equality constraint, so
+    # even though T is invariant, it does imply that (Invariant[T] ≤ Invariant[int]).
+    given_int = ConstraintSet.range(U, T, U) & ConstraintSet.range(int, U, int)
+    static_assert(given_int.implies_subtype_of(Invariant[T], Invariant[int]))
+    static_assert(given_int.implies_subtype_of(Invariant[int], Invariant[T]))
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[bool]))
+    static_assert(not given_int.implies_subtype_of(Invariant[bool], Invariant[T]))
+    static_assert(not given_int.implies_subtype_of(Invariant[T], Invariant[str]))
+    static_assert(not given_int.implies_subtype_of(Invariant[str], Invariant[T]))
+```
+
 [subtyping]: https://typing.python.org/en/latest/spec/concepts.html#subtype-supertype-and-type-equivalence

crates/ty_python_semantic/src/types.rs

@@ -200,7 +200,7 @@ pub(crate) type ApplyTypeMappingVisitor<'db> = TypeTransformer<'db, TypeMapping<
 
 /// A [`PairVisitor`] that is used in `has_relation_to` methods.
 pub(crate) type HasRelationToVisitor<'db> =
-    CycleDetector<TypeRelation, (Type<'db>, Type<'db>, TypeRelation), ConstraintSet<'db>>;
+    CycleDetector<TypeRelation<'db>, (Type<'db>, Type<'db>, TypeRelation<'db>), ConstraintSet<'db>>;
 
 impl Default for HasRelationToVisitor<'_> {
     fn default() -> Self {
@@ -1623,6 +1623,25 @@ impl<'db> Type<'db> {
         self.has_relation_to(db, target, inferable, TypeRelation::Subtyping)
     }
 
+    /// Return the constraints under which this type is a subtype of type `target`, assuming that
+    /// all of the restrictions in `constraints` hold.
+    ///
+    /// See [`TypeRelation::SubtypingAssuming`] for more details.
+    fn when_subtype_of_assuming(
+        self,
+        db: &'db dyn Db,
+        target: Type<'db>,
+        assuming: ConstraintSet<'db>,
+        inferable: InferableTypeVars<'_, 'db>,
+    ) -> ConstraintSet<'db> {
+        self.has_relation_to(
+            db,
+            target,
+            inferable,
+            TypeRelation::SubtypingAssuming(assuming),
+        )
+    }
+
     /// Return true if this type is assignable to type `target`.
     ///
     /// See [`TypeRelation::Assignability`] for more details.
@@ -1654,7 +1673,7 @@ impl<'db> Type<'db> {
         db: &'db dyn Db,
         target: Type<'db>,
         inferable: InferableTypeVars<'_, 'db>,
-        relation: TypeRelation,
+        relation: TypeRelation<'db>,
     ) -> ConstraintSet<'db> {
         self.has_relation_to_impl(
             db,
@@ -1671,7 +1690,7 @@ impl<'db> Type<'db> {
         db: &'db dyn Db,
         target: Type<'db>,
         inferable: InferableTypeVars<'_, 'db>,
-        relation: TypeRelation,
+        relation: TypeRelation<'db>,
         relation_visitor: &HasRelationToVisitor<'db>,
         disjointness_visitor: &IsDisjointVisitor<'db>,
     ) -> ConstraintSet<'db> {
@@ -1684,6 +1703,14 @@ impl<'db> Type<'db> {
             return ConstraintSet::from(true);
         }
 
+        // Handle constraint implication first. If either `self` or `target` is a typevar, check
+        // the constraint set to see if the corresponding constraint is satisfied.
+        if let TypeRelation::SubtypingAssuming(constraints) = relation
+            && (self.is_type_var() || target.is_type_var())
+        {
+            return constraints.implies_subtype_of(db, self, target);
+        }
+
         match (self, target) {
             // Everything is a subtype of `object`.
             (_, Type::NominalInstance(instance)) if instance.is_object() => {
@@ -1763,7 +1790,7 @@ impl<'db> Type<'db> {
                     "DynamicType::Divergent should have been handled in an earlier branch"
                 );
                 ConstraintSet::from(match relation {
-                    TypeRelation::Subtyping => false,
+                    TypeRelation::Subtyping | TypeRelation::SubtypingAssuming(_) => false,
                     TypeRelation::Assignability => true,
                     TypeRelation::Redundancy => match target {
                         Type::Dynamic(_) => true,
@@ -1773,7 +1800,7 @@ impl<'db> Type<'db> {
                 })
             }
             (_, Type::Dynamic(_)) => ConstraintSet::from(match relation {
-                TypeRelation::Subtyping => false,
+                TypeRelation::Subtyping | TypeRelation::SubtypingAssuming(_) => false,
                 TypeRelation::Assignability => true,
                 TypeRelation::Redundancy => match self {
                     Type::Dynamic(_) => true,
@@ -1970,12 +1997,16 @@ impl<'db> Type<'db> {
                     // to non-transitivity (highly undesirable); and pragmatically, a full implementation
                     // of redundancy may not generally lead to simpler types in many situations.
                     let self_ty = match relation {
-                        TypeRelation::Subtyping | TypeRelation::Redundancy => self,
+                        TypeRelation::Subtyping
+                        | TypeRelation::Redundancy
+                        | TypeRelation::SubtypingAssuming(_) => self,
                         TypeRelation::Assignability => self.bottom_materialization(db),
                     };
                     intersection.negative(db).iter().when_all(db, |&neg_ty| {
                         let neg_ty = match relation {
-                            TypeRelation::Subtyping | TypeRelation::Redundancy => neg_ty,
+                            TypeRelation::Subtyping
+                            | TypeRelation::Redundancy
+                            | TypeRelation::SubtypingAssuming(_) => neg_ty,
                             TypeRelation::Assignability => neg_ty.bottom_materialization(db),
                         };
                         self_ty.is_disjoint_from_impl(
@@ -10322,7 +10353,7 @@ impl<'db> ConstructorCallError<'db> {
 
 /// A non-exhaustive enumeration of relations that can exist between types.
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
-pub(crate) enum TypeRelation {
+pub(crate) enum TypeRelation<'db> {
     /// The "subtyping" relation.
     ///
     /// A [fully static] type `B` is a subtype of a fully static type `A` if and only if
@@ -10446,9 +10477,46 @@ pub(crate) enum TypeRelation {
     /// [fully static]: https://typing.python.org/en/latest/spec/glossary.html#term-fully-static-type
     /// [materializations]: https://typing.python.org/en/latest/spec/glossary.html#term-materialize
     Redundancy,
+
+    /// The "constraint implication" relationship, aka "implies subtype of".
+    ///
+    /// This relationship tests whether one type is a [subtype][Self::Subtyping] of another,
+    /// assuming that the constraints in a particular constraint set hold.
+    ///
+    /// For concrete types (types that do not contain typevars), this relationship is the same as
+    /// [subtyping][Self::Subtyping]. (Constraint sets place restrictions on typevars, so if you
+    /// are not comparing typevars, the constraint set can have no effect on whether subtyping
+    /// holds.)
+    ///
+    /// If you're comparing a typevar, we have to consider what restrictions the constraint set
+    /// places on that typevar to determine if subtyping holds. For instance, if you want to check
+    /// whether `T ≤ int`, then the answer will depend on what constraint set you are considering:
+    ///
+    /// ```text
+    /// implies_subtype_of(T ≤ bool, T, int) ⇒ true
+    /// implies_subtype_of(T ≤ int, T, int)  ⇒ true
+    /// implies_subtype_of(T ≤ str, T, int)  ⇒ false
+    /// ```
+    ///
+    /// In the first two cases, the constraint set ensures that `T` will always specialize to a
+    /// type that is a subtype of `int`. In the final case, the constraint set requires `T` to
+    /// specialize to a subtype of `str`, and there is no such type that is also a subtype of
+    /// `int`.
+    ///
+    /// There are two constraint sets that deserve special consideration.
+    ///
+    /// - The "always true" constraint set does not place any restrictions on any typevar. In this
+    ///   case, `implies_subtype_of` will return the same result as `when_subtype_of`, even if
+    ///   you're comparing against a typevar.
+    ///
+    /// - The "always false" constraint set represents an impossible situation. In this case, every
+    ///   subtype check will be vacuously true, even if you're comparing two concrete types that
+    ///   are not actually subtypes of each other. (That is, `implies_subtype_of(false, int, str)`
+    ///   will return true!)
+    SubtypingAssuming(ConstraintSet<'db>),
 }
 
-impl TypeRelation {
+impl TypeRelation<'_> {
     pub(crate) const fn is_assignability(self) -> bool {
         matches!(self, TypeRelation::Assignability)
     }
@@ -10615,7 +10683,7 @@ impl<'db> BoundMethodType<'db> {
         db: &'db dyn Db,
         other: Self,
         inferable: InferableTypeVars<'_, 'db>,
-        relation: TypeRelation,
+        relation: TypeRelation<'db>,
         relation_visitor: &HasRelationToVisitor<'db>,
         disjointness_visitor: &IsDisjointVisitor<'db>,
     ) -> ConstraintSet<'db> {
@@ -10782,7 +10850,7 @@ impl<'db> CallableType<'db> {
         db: &'db dyn Db,
         other: Self,
         inferable: InferableTypeVars<'_, 'db>,
-        relation: TypeRelation,
+        relation: TypeRelation<'db>,
         relation_visitor: &HasRelationToVisitor<'db>,
         disjointness_visitor: &IsDisjointVisitor<'db>,
     ) -> ConstraintSet<'db> {
@@ -10891,7 +10959,7 @@ impl<'db> KnownBoundMethodType<'db> {
         db: &'db dyn Db,
         other: Self,
         inferable: InferableTypeVars<'_, 'db>,
-        relation: TypeRelation,
+        relation: TypeRelation<'db>,
         relation_visitor: &HasRelationToVisitor<'db>,
         disjointness_visitor: &IsDisjointVisitor<'db>,
     ) -> ConstraintSet<'db> {

crates/ty_python_semantic/src/types/call/bind.rs

@@ -1216,10 +1216,10 @@ impl<'db> Bindings<'db> {
                             continue;
                         };
 
-                        let result = tracked.constraints(db).when_subtype_of_given(
+                        let result = ty_a.when_subtype_of_assuming(
                             db,
-                            *ty_a,
                             *ty_b,
+                            tracked.constraints(db),
                             InferableTypeVars::None,
                         );
                         let tracked = TrackedConstraintSet::new(db, result);

crates/ty_python_semantic/src/types/class.rs

@@ -541,14 +541,16 @@ impl<'db> ClassType<'db> {
         db: &'db dyn Db,
         other: Self,
         inferable: InferableTypeVars<'_, 'db>,
-        relation: TypeRelation,
+        relation: TypeRelation<'db>,
         relation_visitor: &HasRelationToVisitor<'db>,
         disjointness_visitor: &IsDisjointVisitor<'db>,
     ) -> ConstraintSet<'db> {
         self.iter_mro(db).when_any(db, |base| {
             match base {
                 ClassBase::Dynamic(_) => match relation {
-                    TypeRelation::Subtyping | TypeRelation::Redundancy => {
+                    TypeRelation::Subtyping
+                    | TypeRelation::Redundancy
+                    | TypeRelation::SubtypingAssuming(_) => {
                         ConstraintSet::from(other.is_object(db))
                     }
                     TypeRelation::Assignability => ConstraintSet::from(!other.is_final(db)),