fix: add types to some functions

Many return and argument types were missing type
annotations.

Author: ishaangandhi

google/cloud/firestore_v1/async_batch.py

@@ -19,6 +19,7 @@
 from google.api_core import retry_async as retries
 
 from google.cloud.firestore_v1.base_batch import BaseWriteBatch
+from google.cloud.firestore_v1.types.write import WriteResult
 
 
 class AsyncWriteBatch(BaseWriteBatch):
@@ -40,7 +41,7 @@ async def commit(
         self,
         retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT,
         timeout: float | None = None,
-    ) -> list:
+    ) -> list[WriteResult]:
         """Commit the changes accumulated in this batch.
 
         Args:

google/cloud/firestore_v1/async_client.py

@@ -25,7 +25,15 @@
 """
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterable, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Iterable,
+    List,
+    Optional,
+    Union,
+)
 
 from google.api_core import gapic_v1
 from google.api_core import retry_async as retries
@@ -50,9 +58,14 @@
 
 if TYPE_CHECKING:  # pragma: NO COVER
     import datetime
+    from typing_extensions import Unpack, TypedDict
 
     from google.cloud.firestore_v1.bulk_writer import BulkWriter
 
+    class TransactionKwargs(TypedDict, total=False):
+        max_attempts: int
+        read_only: bool
+
 
 class AsyncClient(BaseClient):
     """Client for interacting with Google Cloud Firestore API.
@@ -410,7 +423,7 @@ def batch(self) -> AsyncWriteBatch:
         """
         return AsyncWriteBatch(self)
 
-    def transaction(self, **kwargs) -> AsyncTransaction:
+    def transaction(self, **kwargs: Unpack[TransactionKwargs]) -> AsyncTransaction:
         """Get a transaction that uses this client.
 
         See :class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction` for

google/cloud/firestore_v1/base_batch.py

@@ -15,7 +15,7 @@
 """Helpers for batch requests to the Google Cloud Firestore API."""
 from __future__ import annotations
 import abc
-from typing import Dict, Union
+from typing import Any, Dict, Union
 
 # Types needed only for Type Hints
 from google.api_core import retry as retries
@@ -67,7 +67,9 @@ def commit(self):
         write depend on the implementing class."""
         raise NotImplementedError()
 
-    def create(self, reference: BaseDocumentReference, document_data: dict) -> None:
+    def create(
+        self, reference: BaseDocumentReference, document_data: dict[str, Any]
+    ) -> None:
         """Add a "change" to this batch to create a document.
 
         If the document given by ``reference`` already exists, then this
@@ -120,7 +122,7 @@ def set(
     def update(
         self,
         reference: BaseDocumentReference,
-        field_updates: dict,
+        field_updates: dict[str, Any],
         option: _helpers.WriteOption | None = None,
     ) -> None:
         """Add a "change" to update a document.

google/cloud/firestore_v1/base_collection.py

@@ -35,6 +35,7 @@
 from google.api_core import retry as retries
 
 from google.cloud.firestore_v1 import _helpers
+from google.cloud.firestore_v1.base_document import BaseDocumentReference
 from google.cloud.firestore_v1.base_query import QueryType
 
 if TYPE_CHECKING:  # pragma: NO COVER
@@ -133,7 +134,7 @@ def _aggregation_query(self) -> BaseAggregationQuery:
     def _vector_query(self) -> BaseVectorQuery:
         raise NotImplementedError
 
-    def document(self, document_id: Optional[str] = None):
+    def document(self, document_id: Optional[str] = None) -> BaseDocumentReference:
         """Create a sub-document underneath the current collection.
 
         Args:

google/cloud/firestore_v1/base_document.py

@@ -418,7 +418,7 @@ def _client(self):
         return self._reference._client
 
     @property
-    def exists(self):
+    def exists(self) -> bool:
         """Existence flag.
 
         Indicates if the document existed at the time this snapshot
@@ -430,7 +430,7 @@ def exists(self):
         return self._exists
 
     @property
-    def id(self):
+    def id(self) -> str:
         """The document identifier (within its collection).
 
         Returns:
@@ -439,7 +439,7 @@ def id(self):
         return self._reference.id
 
     @property
-    def reference(self):
+    def reference(self) -> BaseDocumentReference:
         """Document reference corresponding to document that owns this data.
 
         Returns:

google/cloud/firestore_v1/document.py

@@ -169,7 +169,7 @@ def set(
 
     def update(
         self,
-        field_updates: dict,
+        field_updates: dict[str, Any],
         option: _helpers.WriteOption | None = None,
         retry: retries.Retry | object | None = gapic_v1.method.DEFAULT,
         timeout: float | None = None,