Implement type equality for recursive types, fixes #161

Author: davidhalter

crates/zuban_python/src/file/inference.rs

@@ -1369,7 +1369,7 @@ impl<'db, 'file> Inference<'db, 'file, '_> {
                     && self
                         .infer_name_of_definition_by_index(first_index)
                         .as_cow_type(i_s)
-                        .is_equal_type(i_s.db, &value.as_cow_type(i_s))
+                        .is_equal_type(i_s.db, None, &value.as_cow_type(i_s))
                 {
                 } else {
                     let mut node_ref = name_def_ref;

crates/zuban_python/src/type_helpers/typing.rs

@@ -7,7 +7,7 @@ use crate::{
     format_data::FormatData,
     inference_state::InferenceState,
     inferred::Inferred,
-    matching::{CouldBeALiteral, Generic, Generics, ResultContext},
+    matching::{CheckedTypeRecursion, CouldBeALiteral, Generic, Generics, ResultContext},
     type_::{
         CallableParams, ClassGenerics, GenericClass, ParamType, StarParamType, StarStarParamType,
         TupleArgs, Type, TypeVarKind, TypedDict, TypedDictGenerics,
@@ -69,7 +69,7 @@ pub(crate) fn execute_cast<'db>(i_s: &InferenceState<'db, '_>, args: &dyn Args<'
     {
         let t_in = actual.as_cow_type(i_s);
         let t_out = result.as_cow_type(i_s);
-        if t_in.is_equal_type(i_s.db, &t_out) && !(t_in.is_any()) {
+        if t_in.is_equal_type(i_s.db, None, &t_out) && !(t_in.is_any()) {
             args.add_issue(
                 i_s,
                 IssueKind::RedundantCast {
@@ -243,7 +243,7 @@ pub(crate) fn execute_assert_type<'db>(
         return Inferred::new_any_from_error();
     };
     let second_type = second.as_cow_type(i_s);
-    if !first_type.is_equal_type(i_s.db, &second_type) {
+    if !first_type.is_equal_type(i_s.db, None, &second_type) {
         let mut format_data = FormatData::new_short(i_s.db);
         format_data.hide_implicit_literals = false;
         let mut actual = first_type.format(&format_data);
@@ -259,8 +259,13 @@ pub(crate) fn execute_assert_type<'db>(
 }
 
 impl Type {
-    pub fn is_equal_type(&self, db: &Database, other: &Type) -> bool {
-        let eq = |t1: &Type, t2: &Type| t1.is_equal_type(db, t2);
+    pub fn is_equal_type(
+        &self,
+        db: &Database,
+        checking_type_recursion: Option<CheckedTypeRecursion>,
+        other: &Type,
+    ) -> bool {
+        let eq = |t1: &Type, t2: &Type| t1.is_equal_type(db, checking_type_recursion, t2);
         let all_eq =
             |ts1: &[Type], ts2: &[Type]| ts1.iter().zip(ts2.iter()).all(|(t1, t2)| eq(t1, t2));
         let typed_dict_eq = |td1: &TypedDict, td2: &TypedDict| {
@@ -276,7 +281,8 @@ impl Type {
                 && match (&m1.extra_items, &m2.extra_items) {
                     (None, None) => true,
                     (Some(t1), Some(t2)) => {
-                        t1.t.is_equal_type(db, &t2.t) && t1.read_only == t2.read_only
+                        t1.t.is_equal_type(db, checking_type_recursion, &t2.t)
+                            && t1.read_only == t2.read_only
                     }
                     _ => false,
                 }
@@ -387,15 +393,16 @@ impl Type {
                     }
             }
             (Type::Type(t1), Type::Type(t2)) => eq(t1, t2),
-            (Type::RecursiveType(r1), Type::RecursiveType(r2)) => {
-                r1.link == r2.link
-                    && r1
-                        .generics
-                        .as_ref()
-                        .zip(r2.generics.as_ref())
-                        .is_none_or(|(g1, g2)| {
-                            matches_generics(Generics::List(g1, None), Generics::List(g2, None))
-                        })
+            (t1 @ Type::RecursiveType(r1), t2) | (t2, t1 @ Type::RecursiveType(r1)) => {
+                let checking_type_recursion = CheckedTypeRecursion {
+                    current: (t1, t2),
+                    previous: checking_type_recursion.as_ref(),
+                };
+                if checking_type_recursion.is_cycle() {
+                    return true;
+                }
+                r1.calculated_type(db)
+                    .is_equal_type(db, Some(checking_type_recursion), t2)
             }
             (Type::Literal(l1), Type::Literal(l2)) => l1.value(db) == l2.value(db),
             (Type::Literal(l), Type::Class(c)) | (Type::Class(c), Type::Literal(l)) => {
@@ -406,12 +413,16 @@ impl Type {
             (Type::Never(_), Type::Never(_)) => true,
             (Type::Union(u1), Type::Union(u2)) => is_equal_union_or_intersection(
                 db,
+                checking_type_recursion,
                 u1.entries.iter().map(|e| &e.type_),
                 u2.entries.iter().map(|e| &e.type_),
             ),
-            (Type::Intersection(i1), Type::Intersection(i2)) => {
-                is_equal_union_or_intersection(db, i1.iter_entries(), i2.iter_entries())
-            }
+            (Type::Intersection(i1), Type::Intersection(i2)) => is_equal_union_or_intersection(
+                db,
+                checking_type_recursion,
+                i1.iter_entries(),
+                i2.iter_entries(),
+            ),
             (Type::EnumMember(m1), Type::EnumMember(m2)) => {
                 m1.member_index == m2.member_index && m1.enum_.defined_at == m2.enum_.defined_at
             }
@@ -431,6 +442,7 @@ impl Type {
 
 fn is_equal_union_or_intersection<'x>(
     db: &Database,
+    checking_type_recursion: Option<CheckedTypeRecursion>,
     ts1: impl ExactSizeIterator<Item = &'x Type>,
     ts2: impl ExactSizeIterator<Item = &'x Type>,
 ) -> bool {
@@ -440,7 +452,7 @@ fn is_equal_union_or_intersection<'x>(
     let mut all_second: Vec<_> = ts2.collect();
     'outer: for t1 in ts1 {
         for (i, t2) in all_second.iter().enumerate() {
-            if t1.is_equal_type(db, t2) {
+            if t1.is_equal_type(db, checking_type_recursion, t2) {
                 all_second.remove(i);
                 continue 'outer;
             }

crates/zuban_python/tests/mypylike/tests/dataclasses.test

@@ -989,3 +989,79 @@ class BaseArithmeticExpression(BaseExpression):
         return Addition()
 
 class Addition[G](BaseArithmeticExpression): ...
+
+[case dataclass_recursive_types_assert_type]
+# From GH #161
+import dataclasses
+from typing import Literal, assert_type, overload
+
+type ArithmeticExpression = int | float | BaseArithmeticExpression
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class BaseExpression[G]:
+    pass
+
+class BaseComparableExpression[G](BaseExpression[G]):
+    pass
+
+class BaseNumericExpression[G](BaseComparableExpression[G]):
+    pass
+
+class BaseArithmeticExpression[G](BaseNumericExpression[G]):
+    @overload
+    def __add__[S: BaseArithmeticExpression[Literal[False]]](
+        self: S,
+        other: int,
+    ) -> Addition[Literal[False], S, int]: ...
+
+    @overload
+    def __add__[S: BaseArithmeticExpression[Literal[False]]](
+        self: S,
+        other: float,
+    ) -> Addition[Literal[False], S, float]: ...
+
+    def __add__(self, other: ArithmeticExpression) -> Addition:
+        return Addition(lhs=self, rhs=other)  # type: ignore
+
+    @overload
+    def __radd__[S: BaseArithmeticExpression[Literal[False]]](
+        self: S,
+        other: int,
+    ) -> Addition[Literal[False], int, S]: ...
+
+    @overload
+    def __radd__[S: BaseArithmeticExpression[Literal[False]]](
+        self: S,
+        other: float,
+    ) -> Addition[Literal[False], float, S]: ...
+
+    def __radd__(self, other: ArithmeticExpression) -> Addition:
+        return Addition(lhs=other, rhs=self)  # type: ignore
+
+class BaseBinaryArithmeticOperation[
+    G,
+    L: ArithmeticExpression,
+    R: ArithmeticExpression,
+](BaseArithmeticExpression[G]):
+    lhs: L
+    rhs: R
+
+class Addition[
+    G,
+    L: ArithmeticExpression,
+    R: ArithmeticExpression,
+](BaseBinaryArithmeticOperation[G, L, R]):
+    pass
+
+class BaseNamedSymbol(BaseExpression[Literal[False]]):
+    name: str
+
+class IntegerNamedSymbol(BaseArithmeticExpression[Literal[False]]):
+    pass
+
+age = IntegerNamedSymbol()
+
+result1 = age + 3.7
+assert_type(result1, Addition[Literal[False], IntegerNamedSymbol, float])
+result2 = 3.7 + age
+assert_type(result2, Addition[Literal[False], float, IntegerNamedSymbol])

crates/zuban_python/tests/mypylike/tests/recursive-type-aliases.test

@@ -200,3 +200,14 @@ IncEx: TypeAlias = Union[
 def foo(a: IncEx | None) -> None: ...
 
 foo(a={"foo"})
+
+[case recursive_type_alias_assert_type]
+from typing import assert_type
+X = list['X'] | int
+
+def f(x: X):
+    assert_type(x, list[X] | int)
+    assert_type(x, X)
+    if isinstance(x, list):
+        assert_type(x[0], list[X] | int)
+        assert_type(x[0], X)