From ca2340513a5ef1cd82d7af0bd8f066f16367887c Mon Sep 17 00:00:00 2001 From: Jillian Kozyra Date: Fri, 9 May 2025 21:59:14 -0700 Subject: [PATCH 1/4] fix: update the async transactional types to not require extra awaits --- .../cloud/firestore_v1/async_transaction.py | 39 ++++++++++++++----- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index 038710929..15327dad5 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/google/cloud/firestore_v1/async_transaction.py @@ -15,7 +15,18 @@ """Helpers for applying Google Cloud Firestore changes in a transaction.""" from __future__ import annotations -from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Coroutine, Optional +from typing import ( + TYPE_CHECKING, + Any, + AsyncGenerator, + Awaitable, + Callable, + Coroutine, + Optional, + TypeVar, + ParamSpec, + Concatenate, +) from google.api_core import exceptions, gapic_v1 from google.api_core import retry_async as retries @@ -41,6 +52,10 @@ from google.cloud.firestore_v1.query_profile import ExplainOptions +T = TypeVar("T") +P = ParamSpec("P") + + class AsyncTransaction(async_batch.AsyncWriteBatch, BaseTransaction): """Accumulate read-and-write operations to be sent in a transaction. @@ -236,11 +251,13 @@ class _AsyncTransactional(_BaseTransactional): A coroutine that should be run (and retried) in a transaction. """ - def __init__(self, to_wrap) -> None: + def __init__( + self, to_wrap: Callable[Concatenate[AsyncTransaction, P], Awaitable[T]] + ) -> None: super(_AsyncTransactional, self).__init__(to_wrap) async def _pre_commit( - self, transaction: AsyncTransaction, *args, **kwargs + self, transaction: AsyncTransaction, *args: P.args, **kwargs: P.kwargs ) -> Coroutine: """Begin transaction and call the wrapped coroutine. @@ -254,7 +271,7 @@ async def _pre_commit( along to the wrapped coroutine. Returns: - Any: result of the wrapped coroutine. + T: result of the wrapped coroutine. Raises: Exception: Any failure caused by ``to_wrap``. @@ -269,12 +286,14 @@ async def _pre_commit( self.retry_id = self.current_id return await self.to_wrap(transaction, *args, **kwargs) - async def __call__(self, transaction, *args, **kwargs): + async def __call__( + self, transaction: AsyncTransaction, *args: P.args, **kwargs: P.kwargs + ) -> T: """Execute the wrapped callable within a transaction. Args: transaction - (:class:`~google.cloud.firestore_v1.transaction.Transaction`): + (:class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction`): A transaction to execute the callable within. args (Tuple[Any, ...]): The extra positional arguments to pass along to the wrapped callable. @@ -282,7 +301,7 @@ async def __call__(self, transaction, *args, **kwargs): along to the wrapped callable. Returns: - Any: The result of the wrapped callable. + T: The result of the wrapped callable. Raises: ValueError: If the transaction does not succeed in @@ -321,13 +340,13 @@ async def __call__(self, transaction, *args, **kwargs): def async_transactional( - to_wrap: Callable[[AsyncTransaction], Any] -) -> _AsyncTransactional: + to_wrap: Callable[Concatenate[AsyncTransaction, P], Awaitable[T]] +) -> Callable[Concatenate[AsyncTransaction, P], Awaitable[T]]: """Decorate a callable so that it runs in a transaction. Args: to_wrap - (Callable[[:class:`~google.cloud.firestore_v1.transaction.Transaction`, ...], Any]): + (Callable[[:class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction`, ...], Any]): A callable that should be run (and retried) in a transaction. Returns: From c0e1446d57e51924fe28b5f93647dc869404f1c2 Mon Sep 17 00:00:00 2001 From: Jillian Kozyra Date: Wed, 21 May 2025 20:51:45 -0700 Subject: [PATCH 2/4] add typing extensions --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 2a47080a1..290eec0e4 100644 --- a/setup.py +++ b/setup.py @@ -39,6 +39,7 @@ "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 = {} From a630701e28c84d58d558881f393d148274a574a6 Mon Sep 17 00:00:00 2001 From: Jillian Kozyra Date: Wed, 21 May 2025 21:10:32 -0700 Subject: [PATCH 3/4] python <3.10 compat --- google/cloud/firestore_v1/async_transaction.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index 15327dad5..416ffaf1e 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/google/cloud/firestore_v1/async_transaction.py @@ -24,10 +24,14 @@ Coroutine, Optional, TypeVar, - ParamSpec, Concatenate, ) +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 From 576a8a8e85d68eac15fb94402422912b4be6ede9 Mon Sep 17 00:00:00 2001 From: Jillian Kozyra Date: Wed, 28 May 2025 22:04:34 -0700 Subject: [PATCH 4/4] use a protocol --- .../cloud/firestore_v1/async_transaction.py | 23 ++++++++----------- setup.py | 1 - 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index 416ffaf1e..b38cdfb71 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/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: diff --git a/setup.py b/setup.py index 290eec0e4..2a47080a1 100644 --- a/setup.py +++ b/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 = {}