From 4a8404107468fb56f85424d288e236408e7fe4cf Mon Sep 17 00:00:00 2001
From: EXPLOSION <git@helvetica.moe>
Date: Tue, 5 Sep 2023 22:12:49 -0400
Subject: [PATCH] Fix KI protection decorator (#2781)

* Fix KI protection decorator

* Appease CI

* Update `verify_types` files

* Try using a TypeVar for callables

* Add `type: ignore`s as necessary

* PR review + CI
---
 trio/_core/_ki.py                     | 20 +++++++++++++-------
 trio/_tests/verify_types_darwin.json  |  2 +-
 trio/_tests/verify_types_linux.json   |  2 +-
 trio/_tests/verify_types_windows.json |  2 +-
 4 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/trio/_core/_ki.py b/trio/_core/_ki.py
index 10172e4989..0ea34619b5 100644
--- a/trio/_core/_ki.py
+++ b/trio/_core/_ki.py
@@ -6,12 +6,13 @@
 import types
 from collections.abc import Callable
 from functools import wraps
-from typing import TYPE_CHECKING, Final, TypeVar
+from typing import TYPE_CHECKING, Final, Protocol, TypeVar
 
 import attr
 
 from .._util import is_main_thread
 
+CallableT = TypeVar("CallableT", bound="Callable[..., object]")
 RetT = TypeVar("RetT")
 
 if TYPE_CHECKING:
@@ -183,14 +184,19 @@ def wrapper(*args: ArgsT.args, **kwargs: ArgsT.kwargs) -> RetT:
     return decorator
 
 
-enable_ki_protection: Callable[
-    [Callable[ArgsT, RetT]], Callable[ArgsT, RetT]
-] = _ki_protection_decorator(True)
+# pyright workaround: https://github.com/microsoft/pyright/issues/5866
+class KIProtectionSignature(Protocol):
+    __name__: str
+
+    def __call__(self, f: CallableT, /) -> CallableT:
+        pass
+
+
+# the following `type: ignore`s are because we use ParamSpec internally, but want to allow overloads
+enable_ki_protection: KIProtectionSignature = _ki_protection_decorator(True)  # type: ignore[assignment]
 enable_ki_protection.__name__ = "enable_ki_protection"
 
-disable_ki_protection: Callable[
-    [Callable[ArgsT, RetT]], Callable[ArgsT, RetT]
-] = _ki_protection_decorator(False)
+disable_ki_protection: KIProtectionSignature = _ki_protection_decorator(False)  # type: ignore[assignment]
 disable_ki_protection.__name__ = "disable_ki_protection"
 
 
diff --git a/trio/_tests/verify_types_darwin.json b/trio/_tests/verify_types_darwin.json
index 2b89d28d8e..3506cb8bee 100644
--- a/trio/_tests/verify_types_darwin.json
+++ b/trio/_tests/verify_types_darwin.json
@@ -76,7 +76,7 @@
     ],
     "otherSymbolCounts": {
       "withAmbiguousType": 0,
-      "withKnownType": 680,
+      "withKnownType": 686,
       "withUnknownType": 0
     },
     "packageName": "trio"
diff --git a/trio/_tests/verify_types_linux.json b/trio/_tests/verify_types_linux.json
index ea5af77abc..640425ea95 100644
--- a/trio/_tests/verify_types_linux.json
+++ b/trio/_tests/verify_types_linux.json
@@ -64,7 +64,7 @@
     ],
     "otherSymbolCounts": {
       "withAmbiguousType": 0,
-      "withKnownType": 680,
+      "withKnownType": 686,
       "withUnknownType": 0
     },
     "packageName": "trio"
diff --git a/trio/_tests/verify_types_windows.json b/trio/_tests/verify_types_windows.json
index 5d3e29a5dc..fff290e159 100644
--- a/trio/_tests/verify_types_windows.json
+++ b/trio/_tests/verify_types_windows.json
@@ -180,7 +180,7 @@
     ],
     "otherSymbolCounts": {
       "withAmbiguousType": 0,
-      "withKnownType": 671,
+      "withKnownType": 677,
       "withUnknownType": 0
     },
     "packageName": "trio"