[ty] Fix panic when trying to pull types for subscript expressions inside Callable type expressions (#18534)

Author: AlexWaygood

crates/ty_project/resources/test/corpus/callable_with_concatenate.py

@@ -0,0 +1,6 @@
+from typing_extensions import TypeVar, Callable, Concatenate, ParamSpec
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+def f(self, callable: Callable[Concatenate[_T, _P], _T]) -> Callable[_P, _T]: ...

crates/ty_project/tests/check.rs

@@ -117,8 +117,10 @@ fn run_corpus_tests(pattern: &str) -> anyhow::Result<()> {
         let code = std::fs::read_to_string(source)?;
 
         let mut check_with_file_name = |path: &SystemPath| {
-            if DO_NOT_ATTEMPT.contains(&&*relative_path.as_str().replace('\\', "/")) {
-                println!("Skipping {relative_path:?} due to known stack overflow");
+            if relative_path.file_name() == Some("types.pyi") {
+                println!(
+                    "Skipping {relative_path:?}: paths with `types.pyi` as their final segment cause a stack overflow"
+                );
                 return;
             }
 
@@ -301,16 +303,4 @@ const KNOWN_FAILURES: &[(&str, bool, bool)] = &[
     // Fails with too-many-cycle-iterations due to a self-referential
     // type alias, see https://github.com/astral-sh/ty/issues/256
     ("crates/ruff_linter/resources/test/fixtures/pyflakes/F401_34.py", true, true),
-
-    // These are all "expression should belong to this TypeInference region and TypeInferenceBuilder should have inferred a type for it"
-    ("crates/ty_vendored/vendor/typeshed/stdlib/abc.pyi", true, true),
-    ("crates/ty_vendored/vendor/typeshed/stdlib/builtins.pyi", true, true),
-    ("crates/ty_vendored/vendor/typeshed/stdlib/curses/__init__.pyi", true, true),
-];
-
-/// Attempting to check one of these files causes a stack overflow
-const DO_NOT_ATTEMPT: &[&str] = &[
-    "crates/ty_vendored/vendor/typeshed/stdlib/pathlib/types.pyi",
-    "crates/ty_vendored/vendor/typeshed/stdlib/types.pyi",
-    "crates/ty_vendored/vendor/typeshed/stdlib/wsgiref/types.pyi",
 ];

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

@@ -252,6 +252,31 @@ def _(c: Callable[[Concatenate[int, str, ...], int], int]):
     reveal_type(c)  # revealed: (...) -> int
 ```
 
+Other type expressions can be nested inside `Concatenate`:
+
+```py
+def _(c: Callable[[Concatenate[int | str, type[str], ...], int], int]):
+    # TODO: Should reveal the correct signature
+    reveal_type(c)  # revealed: (...) -> int
+```
+
+But providing fewer than 2 arguments to `Concatenate` is an error:
+
+```py
+# fmt: off
+
+def _(
+    c: Callable[Concatenate[int], int],  # error: [invalid-type-form] "Special form `typing.Concatenate` expected at least 2 parameters but got 1"
+    d: Callable[Concatenate[(int,)], int],  # error: [invalid-type-form] "Special form `typing.Concatenate` expected at least 2 parameters but got 1"
+    e: Callable[Concatenate[()], int]  # error: [invalid-type-form] "Special form `typing.Concatenate` expected at least 2 parameters but got 0"
+):
+    reveal_type(c)  # revealed: (...) -> int
+    reveal_type(d)  # revealed: (...) -> int
+    reveal_type(e)  # revealed: (...) -> int
+
+# fmt: on
+```
+
 ## Using `typing.ParamSpec`
 
 ```toml

crates/ty_python_semantic/src/types/infer.rs

@@ -9441,8 +9441,29 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                 todo_type!("`TypeGuard[]` special form")
             }
             SpecialFormType::Concatenate => {
-                self.infer_type_expression(arguments_slice);
-                todo_type!("`Concatenate[]` special form")
+                let arguments = if let ast::Expr::Tuple(tuple) = arguments_slice {
+                    &*tuple.elts
+                } else {
+                    std::slice::from_ref(arguments_slice)
+                };
+                for argument in arguments {
+                    self.infer_type_expression(argument);
+                }
+                let num_arguments = arguments.len();
+                let inferred_type = if num_arguments < 2 {
+                    if let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, subscript) {
+                        builder.into_diagnostic(format_args!(
+                            "Special form `{special_form}` expected at least 2 parameters but got {num_arguments}",
+                        ));
+                    }
+                    Type::unknown()
+                } else {
+                    todo_type!("`Concatenate[]` special form")
+                };
+                if arguments_slice.is_tuple_expr() {
+                    self.store_expression_type(arguments_slice, inferred_type);
+                }
+                inferred_type
             }
             SpecialFormType::Unpack => {
                 self.infer_type_expression(arguments_slice);
@@ -9622,7 +9643,9 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                     }))
                 });
             }
-            ast::Expr::Subscript(_) => {
+            ast::Expr::Subscript(subscript) => {
+                let value_ty = self.infer_expression(&subscript.value);
+                self.infer_subscript_type_expression(subscript, value_ty);
                 // TODO: Support `Concatenate[...]`
                 return Some(Parameters::todo());
             }