use a protocol

Author: jekozyra

google/cloud/firestore_v1/async_transaction.py

@@ -24,14 +24,9 @@
     Coroutine,
     Optional,
     TypeVar,
-    Concatenate,
+    Protocol,
 )
 
-try:
-    from typing import ParamSpec
-except ImportError:
-    from typing_extensions import ParamSpec
-
 from google.api_core import exceptions, gapic_v1
 from google.api_core import retry_async as retries
 
@@ -56,8 +51,7 @@
     from google.cloud.firestore_v1.query_profile import ExplainOptions
 
 
-T = TypeVar("T")
-P = ParamSpec("P")
+T = TypeVar("T", bound=Callable[..., Any])
 
 
 class AsyncTransaction(async_batch.AsyncWriteBatch, BaseTransaction):
@@ -256,12 +250,12 @@ class _AsyncTransactional(_BaseTransactional):
     """
 
     def __init__(
-        self, to_wrap: Callable[Concatenate[AsyncTransaction, P], Awaitable[T]]
+        self, to_wrap: Callable[..., Awaitable[T]]
     ) -> None:
         super(_AsyncTransactional, self).__init__(to_wrap)
 
     async def _pre_commit(
-        self, transaction: AsyncTransaction, *args: P.args, **kwargs: P.kwargs
+        self, transaction: AsyncTransaction, *args: Any, **kwargs: Any
     ) -> Coroutine:
         """Begin transaction and call the wrapped coroutine.
 
@@ -291,7 +285,7 @@ async def _pre_commit(
         return await self.to_wrap(transaction, *args, **kwargs)
 
     async def __call__(
-        self, transaction: AsyncTransaction, *args: P.args, **kwargs: P.kwargs
+        self, transaction: AsyncTransaction, *args: Any, **kwargs: Any
     ) -> T:
         """Execute the wrapped callable within a transaction.
 
@@ -343,9 +337,12 @@ async def __call__(
             raise
 
 
+class WithAsyncTransaction(Protocol[T]):
+    def __call__(self, transaction: AsyncTransaction, *args: Any, **kwargs: Any) -> Awaitable[T]: ...
+
 def async_transactional(
-    to_wrap: Callable[Concatenate[AsyncTransaction, P], Awaitable[T]]
-) -> Callable[Concatenate[AsyncTransaction, P], Awaitable[T]]:
+    to_wrap: Callable[..., Awaitable[T]]
+) -> WithAsyncTransaction[T]:
     """Decorate a callable so that it runs in a transaction.
 
     Args:

setup.py

@@ -39,7 +39,6 @@
     "proto-plus >= 1.22.2, <2.0.0; python_version>='3.11'",
     "proto-plus >= 1.25.0, <2.0.0; python_version>='3.13'",
     "protobuf>=3.20.2,<7.0.0dev,!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5",
-    "typing_extensions>=4.13, <5; python_version<'3.10'",
 ]
 extras = {}