[ty] Preserve recursive structure in Annotated values

crates/ty_python_semantic/resources/corpus/recursive_annotated_narrowing.py

@@ -0,0 +1,19 @@
+# Regression test for https://github.com/astral-sh/ty/issues/3837.
+
+from typing import Annotated, TypeIs, reveal_type
+
+
+class Container[T]: ...
+
+
+def is_container[T](value: object, other: T) -> TypeIs[Container[T]]:
+    return True
+
+
+value = int
+while True:
+    if is_container(value, value):
+        value = Annotated[list[value], "metadata"]
+    else:
+        value = {value}
+    reveal_type(value)

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

@@ -172,3 +172,68 @@ class C(Annotated): ...
 
 reveal_mro(C)  # revealed: (<class 'C'>, Unknown, <class 'object'>)
 ```
+
+## Recursive runtime values
+
+`TypeOf` can make an `Annotated` value refer to the type of another runtime `Annotated` object.
+Inference preserves the outer object and collapses the recursive wrapped type. The resulting
+approximation can overlap with a concrete `Annotated` object, so identity comparisons remain
+ambiguous.
+
+```py
+from typing import Annotated
+from ty_extensions._internal import TypeOf
+
+value = int
+while True:
+    value = Annotated[TypeOf[value], "metadata"]
+    reveal_type(value)  # revealed: <special-form 'typing.Annotated[Divergent, <metadata>]'>
+    reveal_type(value is Annotated[TypeOf[int], "metadata"])  # revealed: bool
+    reveal_type(Annotated[TypeOf[int], "metadata"] is value)  # revealed: bool
+    reveal_type(value is not Annotated[TypeOf[int], "metadata"])  # revealed: bool
+```
+
+## Recursive inferred attributes
+
+An inferred attribute can store successive runtime `Annotated` objects. The recursive approximation
+accepts the next assignment, both with and without a separate initializer.
+
+```py
+from typing import Annotated
+from ty_extensions._internal import TypeOf
+
+class Initialized:
+    def __init__(self):
+        self.value = int
+
+    def update(self):
+        self.value = Annotated[TypeOf[self.value], "metadata"]
+
+class Uninitialized:
+    def update(self):
+        self.value = Annotated[TypeOf[self.value], "metadata"]
+```
+
+## Distinct runtime values
+
+An `Annotated` object wrapping `int` is distinct from one wrapping `object`, even though `int` is a
+subtype of `object`.
+
+```py
+from typing import Annotated
+from ty_extensions._internal import TypeOf
+
+integer = Annotated[int, "metadata"]
+object_ = Annotated[object, "metadata"]
+
+reveal_type(integer is object_)  # revealed: Literal[False]
+reveal_type(integer is not object_)  # revealed: Literal[True]
+
+def accepts_integer(value: TypeOf[integer]): ...
+def accepts_object(value: TypeOf[object_]): ...
+
+accepts_integer(integer)
+accepts_object(object_)
+accepts_integer(object_)  # error: [invalid-argument-type]
+accepts_object(integer)  # error: [invalid-argument-type]
+```

crates/ty_python_semantic/src/types/known_instance.rs

@@ -418,9 +418,18 @@ impl<'db> KnownInstanceType<'db> {
             Self::Literal(ty) => ty
                 .recursive_type_normalized_impl(db, env, div, true)
                 .map(Self::Literal),
-            Self::Annotated(ty) => ty
-                .recursive_type_normalized_impl(db, env, div, true)
-                .map(Self::Annotated),
+            Self::Annotated(ty) => {
+                // Keep the runtime Annotated object when its wrapped type diverges. The wrapped
+                // type is still nested: TypeOf can make it another runtime Annotated object.
+                let inner = ty.recursive_type_normalized_impl(db, env, div, true);
+                if nested {
+                    inner.map(Self::Annotated)
+                } else {
+                    Some(Self::Annotated(
+                        inner.unwrap_or_else(|| InternedType::new(db, div)),
+                    ))
+                }
+            }
             Self::TypeGenericAlias(ty) => ty
                 .recursive_type_normalized_impl(db, env, div, true)
                 .map(Self::TypeGenericAlias),

crates/ty_python_semantic/src/types/relation.rs

@@ -1939,6 +1939,24 @@ impl<'a, 'c, 'db> TypeRelationChecker<'a, 'c, 'db> {
                     })
             }
 
+            (
+                Type::KnownInstance(KnownInstanceType::Annotated(source)),
+                Type::KnownInstance(KnownInstanceType::Annotated(target)),
+            ) if source.inner(db).is_recursive_divergent()
+                || target.inner(db).is_recursive_divergent() =>
+            {
+                // Recursive inference preserves an Annotated object while replacing its wrapped
+                // type with Divergent. That approximation must accept another unfolding of the
+                // object, without making distinct concrete Annotated values interchangeable.
+                self.check_relation_in_invariant_position(
+                    db,
+                    source.inner(db),
+                    None,
+                    target.inner(db),
+                    None,
+                )
+            }
+
             // The read-only `__func__` and `__wrapped__` attributes expose the complete wrapped
             // object, so its attributes matter here as well as its call signature.
             (
@@ -3662,6 +3680,17 @@ impl<'a, 'c, 'db> DisjointnessChecker<'a, 'c, 'db> {
                 })
             }),
 
+            (
+                Type::KnownInstance(KnownInstanceType::Annotated(left)),
+                Type::KnownInstance(KnownInstanceType::Annotated(right)),
+            ) if left.inner(db).is_recursive_divergent()
+                || right.inner(db).is_recursive_divergent() =>
+            {
+                // A recursive approximation can represent the same Annotated object as a
+                // concrete unfolding, even though their wrapped types differ.
+                self.never()
+            }
+
             // These types are disjoint whenever their represented objects differ.
             (
                 // `LiteralString` can represent different strings and is handled above.

crates/ty_python_semantic/src/types/tests.rs

@@ -675,6 +675,61 @@ fn pending_narrowing_intersections_are_order_independent() {
     );
 }
 
+#[test]
+fn recursive_annotated_normalization() {
+    let db = setup_db();
+    let env = db.program_environment();
+    let div = Type::divergent(salsa::plumbing::Id::from_bits(1));
+    let list_of = |ty| KnownClass::List.to_specialized_instance(&db, &env, &[ty]);
+    let annotated =
+        |ty| Type::KnownInstance(KnownInstanceType::Annotated(InternedType::new(&db, ty)));
+
+    let recursive = annotated(div);
+    assert_eq!(
+        annotated(list_of(div)).recursive_type_normalized_impl(&db, &env, div, false),
+        Some(recursive)
+    );
+    assert_eq!(
+        annotated(list_of(list_of(div))).recursive_type_normalized_impl(&db, &env, div, false),
+        Some(recursive)
+    );
+    assert_eq!(
+        annotated(recursive).recursive_type_normalized_impl(&db, &env, div, false),
+        Some(recursive)
+    );
+    assert_eq!(
+        annotated(list_of(div)).recursive_type_normalized_impl(&db, &env, div, true),
+        None
+    );
+}
+
+#[test]
+fn recursive_annotated_relations() {
+    let db = setup_db();
+    let env = db.program_environment();
+    let div = Type::divergent(salsa::plumbing::Id::from_bits(1));
+    let annotated =
+        |ty| Type::KnownInstance(KnownInstanceType::Annotated(InternedType::new(&db, ty)));
+    let recursive = annotated(div);
+    let integer = annotated(KnownClass::Int.to_instance(&db, &env));
+    let object = annotated(Type::object());
+
+    for concrete in [integer, object, annotated(recursive)] {
+        assert!(concrete.is_assignable_to(&db, &env, recursive));
+        assert!(recursive.is_assignable_to(&db, &env, concrete));
+        assert!(!concrete.is_subtype_of(&db, &env, recursive));
+        assert!(!recursive.is_subtype_of(&db, &env, concrete));
+        assert!(!concrete.is_disjoint_from(&db, &env, recursive));
+        assert!(!recursive.is_disjoint_from(&db, &env, concrete));
+    }
+    assert!(!integer.is_assignable_to(&db, &env, object));
+    assert!(!object.is_assignable_to(&db, &env, integer));
+    assert!(integer.is_disjoint_from(&db, &env, object));
+    assert!(object.is_disjoint_from(&db, &env, integer));
+    assert!(!Type::int_literal(1).is_assignable_to(&db, &env, recursive));
+    assert!(!recursive.is_assignable_to(&db, &env, KnownClass::Int.to_instance(&db, &env)));
+}
+
 #[test]
 fn pending_narrowing_cycle_recovery_preserves_guarded_contributions() {
     let db = setup_db();