Fix KI protection decorator (python-trio#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
apollo13 · Sep 6, 2023 · 4a84041 · 4a84041
1 parent 005d41d
commit 4a84041
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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
@@ -180,7 +180,7 @@
     ],
     "otherSymbolCounts": {
       "withAmbiguousType": 0,
-      "withKnownType": 671,
+      "withKnownType": 677,
       "withUnknownType": 0
     },
     "packageName": "trio"