diff --git a/google/cloud/firestore_v1/async_aggregation.py b/google/cloud/firestore_v1/async_aggregation.py index c39b50c5e4..5bab339c57 100644 --- a/google/cloud/firestore_v1/async_aggregation.py +++ b/google/cloud/firestore_v1/async_aggregation.py @@ -23,9 +23,9 @@ from google.api_core import gapic_v1 from google.api_core import retry_async as retries -from typing import List, Union, AsyncGenerator - +from typing import List, Union, AsyncGenerator, Optional +from google.cloud.firestore_v1.async_transaction import AsyncTransaction from google.cloud.firestore_v1.base_aggregation import ( AggregationResult, _query_response_to_result, @@ -44,7 +44,7 @@ def __init__( async def get( self, - transaction=None, + transaction: Optional[AsyncTransaction] = None, retry: Union[ retries.AsyncRetry, None, gapic_v1.method._MethodDefault ] = gapic_v1.method.DEFAULT, @@ -78,7 +78,7 @@ async def get( async def stream( self, - transaction=None, + transaction: Optional[AsyncTransaction] = None, retry: Union[ retries.AsyncRetry, None, gapic_v1.method._MethodDefault ] = gapic_v1.method.DEFAULT, diff --git a/google/cloud/firestore_v1/async_batch.py b/google/cloud/firestore_v1/async_batch.py index 84b45fa094..64324cc3cd 100644 --- a/google/cloud/firestore_v1/async_batch.py +++ b/google/cloud/firestore_v1/async_batch.py @@ -19,6 +19,8 @@ from google.api_core import retry_async as retries from google.cloud.firestore_v1.base_batch import BaseWriteBatch +from google.cloud.firestore_v1.types import WriteResult +from typing import Optional class AsyncWriteBatch(BaseWriteBatch): @@ -39,8 +41,8 @@ def __init__(self, client) -> None: async def commit( self, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, - ) -> list: + timeout: Optional[float] = None, + ) -> list[WriteResult]: """Commit the changes accumulated in this batch. Args: diff --git a/google/cloud/firestore_v1/async_client.py b/google/cloud/firestore_v1/async_client.py index 20541c3770..8e163d915b 100644 --- a/google/cloud/firestore_v1/async_client.py +++ b/google/cloud/firestore_v1/async_client.py @@ -226,10 +226,10 @@ def document(self, *document_path: str) -> AsyncDocumentReference: async def get_all( self, references: List[AsyncDocumentReference], - field_paths: Iterable[str] = None, - transaction=None, + field_paths: Optional[Iterable[str]] = None, + transaction: Optional[AsyncTransaction] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncGenerator[DocumentSnapshot, Any]: """Retrieve a batch of documents. @@ -285,7 +285,7 @@ async def get_all( async def collections( self, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncGenerator[AsyncCollectionReference, Any]: """List top-level collections of the client's database. diff --git a/google/cloud/firestore_v1/async_collection.py b/google/cloud/firestore_v1/async_collection.py index 093117d40b..5a83b65ddc 100644 --- a/google/cloud/firestore_v1/async_collection.py +++ b/google/cloud/firestore_v1/async_collection.py @@ -22,14 +22,13 @@ _item_to_document_ref, ) from google.cloud.firestore_v1 import async_query, async_document, async_aggregation +from google.cloud.firestore_v1.async_transaction import AsyncTransaction +from google.protobuf.timestamp_pb2 import Timestamp from google.cloud.firestore_v1.document import DocumentReference from typing import AsyncIterator -from typing import Any, AsyncGenerator, Tuple - -# Types needed only for Type Hints -from google.cloud.firestore_v1.transaction import Transaction +from typing import AsyncGenerator, Tuple, Optional class AsyncCollectionReference(BaseCollectionReference[async_query.AsyncQuery]): @@ -84,10 +83,10 @@ async def _chunkify(self, chunk_size: int): async def add( self, document_data: dict, - document_id: str = None, + document_id: Optional[str] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, - ) -> Tuple[Any, Any]: + timeout: Optional[float] = None, + ) -> Tuple[Timestamp, async_document.AsyncDocumentReference]: """Create a document in the Firestore database with the provided data. Args: @@ -125,7 +124,7 @@ async def add( return write_result.update_time, document_ref def document( - self, document_id: str = None + self, document_id: Optional[str] = None ) -> async_document.AsyncDocumentReference: """Create a sub-document underneath the current collection. @@ -143,9 +142,9 @@ def document( async def list_documents( self, - page_size: int = None, + page_size: Optional[int] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncGenerator[DocumentReference, None]: """List all subdocuments of the current collection. @@ -176,9 +175,9 @@ async def list_documents( async def get( self, - transaction: Transaction = None, + transaction: Optional[AsyncTransaction] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> list: """Read the documents in this collection. @@ -207,9 +206,9 @@ async def get( async def stream( self, - transaction: Transaction = None, + transaction: Optional[AsyncTransaction] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncIterator[async_document.DocumentSnapshot]: """Read the documents in this collection. diff --git a/google/cloud/firestore_v1/async_document.py b/google/cloud/firestore_v1/async_document.py index 75250d0b4c..a110b78c66 100644 --- a/google/cloud/firestore_v1/async_document.py +++ b/google/cloud/firestore_v1/async_document.py @@ -26,9 +26,10 @@ _first_write_result, ) from google.cloud.firestore_v1 import _helpers +from google.cloud.firestore_v1.async_transaction import AsyncTransaction from google.cloud.firestore_v1.types import write from google.protobuf.timestamp_pb2 import Timestamp -from typing import AsyncGenerator, Iterable +from typing import AsyncGenerator, Iterable, Optional logger = logging.getLogger(__name__) @@ -66,7 +67,7 @@ async def create( self, document_data: dict, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> write.WriteResult: """Create the current document in the Firestore database. @@ -96,7 +97,7 @@ async def set( document_data: dict, merge: bool = False, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> write.WriteResult: """Replace the current document in the Firestore database. @@ -134,9 +135,9 @@ async def set( async def update( self, field_updates: dict, - option: _helpers.WriteOption = None, + option: Optional[_helpers.WriteOption] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> write.WriteResult: """Update an existing document in the Firestore database. @@ -291,9 +292,9 @@ async def update( async def delete( self, - option: _helpers.WriteOption = None, + option: Optional[_helpers.WriteOption] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> Timestamp: """Delete the current document in the Firestore database. @@ -325,10 +326,10 @@ async def delete( async def get( self, - field_paths: Iterable[str] = None, - transaction=None, + field_paths: Optional[Iterable[str]] = None, + transaction: Optional[AsyncTransaction] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> DocumentSnapshot: """Retrieve a snapshot of the current document. @@ -394,9 +395,9 @@ async def get( async def collections( self, - page_size: int = None, + page_size: Optional[int] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncGenerator: """List subcollections of the current document. diff --git a/google/cloud/firestore_v1/async_query.py b/google/cloud/firestore_v1/async_query.py index 8ee4012904..4af4b68455 100644 --- a/google/cloud/firestore_v1/async_query.py +++ b/google/cloud/firestore_v1/async_query.py @@ -40,7 +40,7 @@ if TYPE_CHECKING: # pragma: NO COVER # Types needed only for Type Hints - from google.cloud.firestore_v1.transaction import Transaction + from google.cloud.firestore_v1.async_transaction import AsyncTransaction from google.cloud.firestore_v1.field_path import FieldPath @@ -171,9 +171,9 @@ async def _chunkify( async def get( self, - transaction: Transaction = None, + transaction: Optional[AsyncTransaction] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> list: """Read the documents in the collection that match this query. @@ -266,9 +266,9 @@ def avg( async def stream( self, - transaction=None, + transaction: Optional[AsyncTransaction] = None, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncGenerator[async_document.DocumentSnapshot, None]: """Read the documents in the collection that match this query. @@ -379,9 +379,9 @@ def _get_query_class(): async def get_partitions( self, - partition_count, + partition_count: int, retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncGenerator[QueryPartition, None]: """Partition a query for parallelization. diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index 18a20b8e12..6309290521 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/google/cloud/firestore_v1/async_transaction.py @@ -43,7 +43,7 @@ from google.cloud.firestore_v1.async_document import AsyncDocumentReference from google.cloud.firestore_v1.async_document import DocumentSnapshot from google.cloud.firestore_v1.async_query import AsyncQuery -from typing import Any, AsyncGenerator, Callable, Coroutine +from typing import Any, AsyncGenerator, Callable, Coroutine, Optional, Union # Types needed only for Type Hints from google.cloud.firestore_v1.client import Client @@ -82,7 +82,7 @@ def _add_write_pbs(self, write_pbs: list) -> None: super(AsyncTransaction, self)._add_write_pbs(write_pbs) - async def _begin(self, retry_id: bytes = None) -> None: + async def _begin(self, retry_id: Optional[bytes] = None) -> None: """Begin the transaction. Args: @@ -152,9 +152,9 @@ async def _commit(self) -> list: async def get_all( self, - references: list, + references: list[AsyncDocumentReference], retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncGenerator[DocumentSnapshot, Any]: """Retrieves multiple documents from Firestore. @@ -175,9 +175,9 @@ async def get_all( async def get( self, - ref_or_query, + ref_or_query: Union[AsyncDocumentReference, AsyncQuery], retry: retries.AsyncRetry = gapic_v1.method.DEFAULT, - timeout: float = None, + timeout: Optional[float] = None, ) -> AsyncGenerator[DocumentSnapshot, Any]: """ Retrieve a document or a query result from the database.