fix(types): improve typing (#1136)

A few types were missing, which makes it hard to use Firestore with
strict typechecking.

Author: ishaangandhi

google/cloud/firestore_v1/async_batch.py

@@ -19,6 +19,7 @@
 fromgoogle.api_coreimportretry_asyncasretries
 
 fromgoogle.cloud.firestore_v1.base_batchimportBaseWriteBatch
+fromgoogle.cloud.firestore_v1.types.writeimportWriteResult
 
 
 classAsyncWriteBatch(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__ importannotations
 
-fromtypingimportTYPE_CHECKING, Any, AsyncGenerator, Iterable, List, Optional, Union
+fromtypingimport (
+TYPE_CHECKING,
+Any,
+AsyncGenerator,
+Iterable,
+List,
+Optional,
+Union,
+)
 
 fromgoogle.api_coreimportgapic_v1
 fromgoogle.api_coreimportretry_asyncasretries
@@ -40,6 +48,7 @@
 fromgoogle.cloud.firestore_v1.async_transactionimportAsyncTransaction
 fromgoogle.cloud.firestore_v1.base_clientimport_parse_batch_get# type: ignore
 fromgoogle.cloud.firestore_v1.base_clientimport_CLIENT_INFO, BaseClient, _path_helper
+fromgoogle.cloud.firestore_v1.base_transactionimportMAX_ATTEMPTS
 fromgoogle.cloud.firestore_v1.field_pathimportFieldPath
 fromgoogle.cloud.firestore_v1.services.firestoreimport (
 async_clientasfirestore_client,
@@ -410,7 +419,9 @@ def batch(self) -> AsyncWriteBatch:
  """
 returnAsyncWriteBatch(self)
 
-deftransaction(self, **kwargs) ->AsyncTransaction:
+deftransaction(
+self, max_attempts: int=MAX_ATTEMPTS, read_only: bool=False
+ ) ->AsyncTransaction:
 """Get a transaction that uses this client.
 
  See :class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction` for
@@ -426,4 +437,4 @@ def transaction(self, **kwargs) -> AsyncTransaction:
  :class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction`:
  A transaction attached to this client.
  """
-returnAsyncTransaction(self, **kwargs)
+returnAsyncTransaction(self, max_attempts=max_attempts, read_only=read_only)

google/cloud/firestore_v1/async_collection.py

@@ -15,7 +15,7 @@
 """Classes for representing collections for the Google Cloud Firestore API."""
 from __future__ importannotations
 
-fromtypingimportTYPE_CHECKING, Any, AsyncGenerator, Optional, Tuple
+fromtypingimportTYPE_CHECKING, Any, AsyncGenerator, Optional, Tuple, cast
 
 fromgoogle.api_coreimportgapic_v1
 fromgoogle.api_coreimportretry_asyncasretries
@@ -153,7 +153,8 @@ def document(self, document_id: str | None = None) -> AsyncDocumentReference:
  :class:`~google.cloud.firestore_v1.document.async_document.AsyncDocumentReference`:
  The child document.
  """
-returnsuper(AsyncCollectionReference, self).document(document_id)
+doc=super(AsyncCollectionReference, self).document(document_id)
+returncast("AsyncDocumentReference", doc)
 
 asyncdeflist_documents(
 self,

google/cloud/firestore_v1/base_batch.py

@@ -15,7 +15,7 @@
 """Helpers for batch requests to the Google Cloud Firestore API."""
 from __future__ importannotations
 importabc
-fromtypingimportDict, Union
+fromtypingimportAny, Dict, Union
 
 # Types needed only for Type Hints
 fromgoogle.api_coreimportretryasretries
@@ -67,7 +67,9 @@ def commit(self):
  write depend on the implementing class."""
 raiseNotImplementedError()
 
-defcreate(self, reference: BaseDocumentReference, document_data: dict) ->None:
+defcreate(
+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(
 defupdate(
 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_client.py

@@ -57,7 +57,7 @@
 DocumentSnapshot,
 )
 fromgoogle.cloud.firestore_v1.base_queryimportBaseQuery
-fromgoogle.cloud.firestore_v1.base_transactionimportBaseTransaction
+fromgoogle.cloud.firestore_v1.base_transactionimportMAX_ATTEMPTS, BaseTransaction
 fromgoogle.cloud.firestore_v1.bulk_writerimportBulkWriter, BulkWriterOptions
 fromgoogle.cloud.firestore_v1.field_pathimportrender_field_path
 fromgoogle.cloud.firestore_v1.services.firestoreimportclientasfirestore_client
@@ -497,7 +497,9 @@ def collections(
 defbatch(self) ->BaseWriteBatch:
 raiseNotImplementedError
 
-deftransaction(self, **kwargs) ->BaseTransaction:
+deftransaction(
+self, max_attempts: int=MAX_ATTEMPTS, read_only: bool=False
+ ) ->BaseTransaction:
 raiseNotImplementedError
 
 

google/cloud/firestore_v1/base_collection.py

@@ -35,6 +35,7 @@
 fromgoogle.api_coreimportretryasretries
 
 fromgoogle.cloud.firestore_v1import_helpers
+fromgoogle.cloud.firestore_v1.base_documentimportBaseDocumentReference
 fromgoogle.cloud.firestore_v1.base_queryimportQueryType
 
 ifTYPE_CHECKING: # pragma: NO COVER
@@ -133,7 +134,7 @@ def _aggregation_query(self) -> BaseAggregationQuery:
 def_vector_query(self) ->BaseVectorQuery:
 raiseNotImplementedError
 
-defdocument(self, document_id: Optional[str] =None):
+defdocument(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):
 returnself._reference._client
 
 @property
-defexists(self):
+defexists(self)->bool:
 """Existence flag.
 
  Indicates if the document existed at the time this snapshot
@@ -430,7 +430,7 @@ def exists(self):
 returnself._exists
 
 @property
-defid(self):
+defid(self)->str:
 """The document identifier (within its collection).
 
  Returns:
@@ -439,7 +439,7 @@ def id(self):
 returnself._reference.id
 
 @property
-defreference(self):
+defreference(self)->BaseDocumentReference:
 """Document reference corresponding to document that owns this data.
 
  Returns:

google/cloud/firestore_v1/client.py

@@ -39,6 +39,7 @@
 
 # Types needed only for Type Hints
 fromgoogle.cloud.firestore_v1.base_documentimportDocumentSnapshot
+fromgoogle.cloud.firestore_v1.base_transactionimportMAX_ATTEMPTS
 fromgoogle.cloud.firestore_v1.batchimportWriteBatch
 fromgoogle.cloud.firestore_v1.collectionimportCollectionReference
 fromgoogle.cloud.firestore_v1.documentimportDocumentReference
@@ -391,7 +392,9 @@ def batch(self) -> WriteBatch:
  """
 returnWriteBatch(self)
 
-deftransaction(self, **kwargs) ->Transaction:
+deftransaction(
+self, max_attempts: int=MAX_ATTEMPTS, read_only: bool=False
+ ) ->Transaction:
 """Get a transaction that uses this client.
 
  See :class:`~google.cloud.firestore_v1.transaction.Transaction` for
@@ -407,4 +410,4 @@ def transaction(self, **kwargs) -> Transaction:
  :class:`~google.cloud.firestore_v1.transaction.Transaction`:
  A transaction attached to this client.
  """
-returnTransaction(self, **kwargs)
+returnTransaction(self, max_attempts=max_attempts, read_only=read_only)

google/cloud/firestore_v1/document.py

@@ -169,7 +169,7 @@ def set(
 
 defupdate(
 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,