From 4e68dd96a61c4fe0411d4dff832d4c60084349fb Mon Sep 17 00:00:00 2001
From: David Peter <sharkdp@users.noreply.github.com>
Date: Mon, 26 May 2025 17:08:52 +0200
Subject: [PATCH] [ty] Infer types for ty_extensions.Intersection[A, B] tuple
 expressions (#18321)

## Summary

fixes astral-sh/ty#366

## Test Plan

* Added panic corpus regression tests
* I also wrote a hover regression test (see below), but decided not to
include it. The corpus tests are much more "effective" at finding these
types of errors, since they exhaustively check all expressions for
types.

<details>

```rs
#[test]
fn hover_regression_test_366() {
    let test = cursor_test(
        r#"
    from ty_extensions import Intersection

    class A: ...
    class B: ...

    def _(x: Intersection[A,<CURSOR> B]):
        pass
    "#,
    );

    assert_snapshot!(test.hover(), @r"
    A & B
    ---------------------------------------------
    ```text
    A & B
    ```
    ---------------------------------------------
    info[hover]: Hovered content is
     --> main.py:7:31
      |
    5 |         class B: ...
    6 |
    7 |         def _(x: Intersection[A, B]):
      |                               ^^-^
      |                               | |
      |                               | Cursor offset
      |                               source
    8 |             pass
      |
    ");
}
```

</details>
---
 .../resources/test/corpus/ty_extensions.py    | 30 +++++++++++++++++++
 crates/ty_python_semantic/src/types/infer.rs  |  9 ++++--
 2 files changed, 37 insertions(+), 2 deletions(-)
 create mode 100644 crates/ty_project/resources/test/corpus/ty_extensions.py

diff --git a/crates/ty_project/resources/test/corpus/ty_extensions.py b/crates/ty_project/resources/test/corpus/ty_extensions.py
new file mode 100644
index 0000000000..64810a6307
--- /dev/null
+++ b/crates/ty_project/resources/test/corpus/ty_extensions.py
@@ -0,0 +1,30 @@
+"""
+Make sure that types are inferred for all subexpressions of the following
+annotations involving ty_extension `_SpecialForm`s.
+
+This is a regression test for https://github.com/astral-sh/ty/issues/366
+"""
+
+from ty_extensions import CallableTypeOf, Intersection, Not, TypeOf
+
+
+class A: ...
+
+
+class B: ...
+
+
+def _(x: Not[A]):
+    pass
+
+
+def _(x: Intersection[A], y: Intersection[A, B]):
+    pass
+
+
+def _(x: TypeOf[1j]):
+    pass
+
+
+def _(x: CallableTypeOf[str]):
+    pass
diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs
index bddbe9f836..b60ac4c9ba 100644
--- a/crates/ty_python_semantic/src/types/infer.rs
+++ b/crates/ty_python_semantic/src/types/infer.rs
@@ -8648,11 +8648,16 @@ impl<'db> TypeInferenceBuilder<'db> {
                     element => Either::Right(std::iter::once(element)),
                 };
 
-                elements
+                let ty = elements
                     .fold(IntersectionBuilder::new(db), |builder, element| {
                         builder.add_positive(self.infer_type_expression(element))
                     })
-                    .build()
+                    .build();
+
+                if matches!(arguments_slice, ast::Expr::Tuple(_)) {
+                    self.store_expression_type(arguments_slice, ty);
+                }
+                ty
             }
             KnownInstanceType::TypeOf => match arguments_slice {
                 ast::Expr::Tuple(_) => {