From 2190fcebe0eb3d5d7160e23f0392c2d28a172635 Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Wed, 7 Jan 2026 13:30:34 +0530
Subject: [PATCH] [ty] Substitute `ParamSpec` in overloaded functions (#22416)

## Summary

fixes: https://github.com/astral-sh/ty/issues/2027

This PR fixes a bug where the type mapping for a `ParamSpec` was not
being applied in an overloaded function.

This PR also fixes https://github.com/astral-sh/ty/issues/2081 and
reveals new diagnostics which doesn't look related to the bug:

```py
from prefect import flow, task

@task
def task_get() -> int:
    """Task get integer."""
    return 42

@task
def task_add(x: int, y: int) -> int:
    """Task add two integers."""
    print(f"Adding {x} and {y}")
    return x + y

@flow
def my_flow():
    """My flow."""
    x = 23
    future_y = task_get.submit()

	# error: [no-matching-overload]
    task_add(future_y, future_y)
	# error: [no-matching-overload]
    task_add(x, future_y)
```

The reason is that the type of `future_y` is `PrefectFuture[int]` while
the type of `task_add` is `Task[(x: int, y: int), int]` which means that
the assignment between `int` and `PrefectFuture[int]` fails which
results in no overload matching. Pyright also raises the invalid
argument type error on all three usages of `future_y` in those two
calls.

## Test Plan

Add regression mdtest from the linked issue.
---
 .../mdtest/generics/pep695/paramspec.md       | 40 ++++++++++++++++++
 .../src/types/signatures.rs                   | 42 +++++++++++--------
 2 files changed, 65 insertions(+), 17 deletions(-)

diff --git a/crates/ty_python_semantic/resources/mdtest/generics/pep695/paramspec.md b/crates/ty_python_semantic/resources/mdtest/generics/pep695/paramspec.md
index 0aee44902e..03883122ee 100644
--- a/crates/ty_python_semantic/resources/mdtest/generics/pep695/paramspec.md
+++ b/crates/ty_python_semantic/resources/mdtest/generics/pep695/paramspec.md
@@ -688,6 +688,46 @@ reveal_type(with_parameters(int_int, 1))  # revealed: Overload[(x: int) -> str,
 reveal_type(with_parameters(int_int, "a"))  # revealed: Overload[(x: int) -> str, (x: str) -> str]
 ```
 
+### Overloads with subtitution of `P.args` and `P.kwargs`
+
+This is regression test for <https://github.com/astral-sh/ty/issues/2027>
+
+```py
+from typing import Callable, Never, overload
+
+class Task[**P, R]:
+    def __init__(self, func: Callable[P, R]) -> None:
+        self.func = func
+
+    @overload
+    def __call__(self: "Task[P, R]", *args: P.args, **kwargs: P.kwargs) -> R: ...
+    @overload
+    def __call__(self: "Task[P, Never]", *args: P.args, **kwargs: P.kwargs) -> None: ...
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R | None:
+        return self.func(*args, **kwargs)
+
+def returns_str(x: int) -> str:
+    return str(x)
+
+def never_returns(x: int) -> Never:
+    raise Exception()
+
+t1 = Task(returns_str)
+reveal_type(t1)  # revealed: Task[(x: int), str]
+reveal_type(t1(1))  # revealed: str
+reveal_type(t1(x=1))  # revealed: str
+# error: [no-matching-overload]
+reveal_type(t1("a"))  # revealed: Unknown
+# error: [no-matching-overload]
+reveal_type(t1(y=1))  # revealed: Unknown
+
+t2 = Task(never_returns)
+# TODO: This should be `Task[(x: int), Never]`
+reveal_type(t2)  # revealed: Task[(x: int), Unknown]
+# TODO: This should be `Never`
+reveal_type(t2(1))  # revealed: Unknown
+```
+
 ## ParamSpec attribute assignability
 
 When comparing signatures with `ParamSpec` attributes (`P.args` and `P.kwargs`), two different
diff --git a/crates/ty_python_semantic/src/types/signatures.rs b/crates/ty_python_semantic/src/types/signatures.rs
index bf6dc2b3b9..a58bf764c5 100644
--- a/crates/ty_python_semantic/src/types/signatures.rs
+++ b/crates/ty_python_semantic/src/types/signatures.rs
@@ -219,23 +219,31 @@ impl<'db> CallableSignature<'db> {
             }
         }
 
-        if let TypeMapping::ApplySpecialization(specialization) = type_mapping
-            && let [self_signature] = self.overloads.as_slice()
-            && let Some((prefix_parameters, paramspec)) = self_signature
-                .parameters
-                .find_paramspec_from_args_kwargs(db)
-            && let Some(paramspec_value) = specialization.get(db, paramspec)
-            && let Some(result) = try_apply_type_mapping_for_paramspec(
-                db,
-                self_signature,
-                prefix_parameters,
-                paramspec_value,
-                type_mapping,
-                tcx,
-                visitor,
-            )
-        {
-            result
+        if let TypeMapping::ApplySpecialization(specialization) = type_mapping {
+            Self::from_overloads(self.overloads.iter().flat_map(|signature| {
+                if let Some((prefix, paramspec)) =
+                    signature.parameters.find_paramspec_from_args_kwargs(db)
+                    && let Some(value) = specialization.get(db, paramspec)
+                    && let Some(result) = try_apply_type_mapping_for_paramspec(
+                        db,
+                        signature,
+                        prefix,
+                        value,
+                        type_mapping,
+                        tcx,
+                        visitor,
+                    )
+                {
+                    result.overloads
+                } else {
+                    smallvec_inline![signature.apply_type_mapping_impl(
+                        db,
+                        type_mapping,
+                        tcx,
+                        visitor
+                    )]
+                }
+            }))
         } else {
             Self::from_overloads(
                 self.overloads.iter().map(|signature| {