feat: improve type information

google/cloud/firestore_v1/async_document.py

@@ -25,7 +25,7 @@
 from google.cloud.firestore_v1.types import common
 from google.cloud.firestore_v1.types import write
 from google.protobuf import timestamp_pb2
-from typing import Any, AsyncGenerator, Coroutine, Union
+from typing import Any, AsyncGenerator, Coroutine, Iterable, Union
 
 
 class AsyncDocumentReference(BaseDocumentReference):
@@ -56,7 +56,7 @@ class AsyncDocumentReference(BaseDocumentReference):
     def __init__(self, *path, **kwargs) -> None:
         super(AsyncDocumentReference, self).__init__(*path, **kwargs)
 
-    async def create(self, document_data) -> write.WriteResult:
+    async def create(self, document_data: dict) -> write.WriteResult:
         """Create the current document in the Firestore database.
 
         Args:
@@ -77,7 +77,7 @@ async def create(self, document_data) -> write.WriteResult:
         write_results = await batch.commit()
         return _first_write_result(write_results)
 
-    async def set(self, document_data, merge=False) -> write.WriteResult:
+    async def set(self, document_data: dict, merge: bool = False) -> write.WriteResult:
         """Replace the current document in the Firestore database.
 
         A write ``option`` can be specified to indicate preconditions of
@@ -108,7 +108,9 @@ async def set(self, document_data, merge=False) -> write.WriteResult:
         write_results = await batch.commit()
         return _first_write_result(write_results)
 
-    async def update(self, field_updates, option=None) -> write.WriteResult:
+    async def update(
+        self, field_updates: dict, option: _helpers.WriteOption = None
+    ) -> write.WriteResult:
         """Update an existing document in the Firestore database.
 
         By default, this method verifies that the document exists on the
@@ -256,7 +258,9 @@ async def update(self, field_updates, option=None) -> write.WriteResult:
         write_results = await batch.commit()
         return _first_write_result(write_results)
 
-    async def delete(self, option=None) -> timestamp_pb2.Timestamp:
+    async def delete(
+        self, option: _helpers.WriteOption = None
+    ) -> timestamp_pb2.Timestamp:
         """Delete the current document in the Firestore database.
 
         Args:

google/cloud/firestore_v1/async_transaction.py

@@ -144,7 +144,7 @@ async def _commit(self) -> list:
         self._clean_up()
         return list(commit_response.write_results)
 
-    async def get_all(self, references) -> AsyncGenerator[DocumentSnapshot, Any]:
+    async def get_all(self, references: list) -> AsyncGenerator[DocumentSnapshot, Any]:
         """Retrieves multiple documents from Firestore.
 
         Args:

google/cloud/firestore_v1/base_client.py

@@ -272,7 +272,7 @@ def _document_path_helper(self, *document_path) -> List[str]:
         return joined_path.split(_helpers.DOCUMENT_PATH_DELIMITER)
 
     @staticmethod
-    def field_path(*field_names) -> str:
+    def field_path(*field_names: Tuple[str]) -> str:
         """Create a **field path** from a list of nested field names.
 
         A **field path** is a ``.``-delimited concatenation of the field
@@ -405,7 +405,7 @@ def _reference_info(references: list) -> Tuple[list, dict]:
     return document_paths, reference_map
 
 
-def _get_reference(document_path, reference_map) -> BaseDocumentReference:
+def _get_reference(document_path: str, reference_map: dict) -> BaseDocumentReference:
     """Get a document reference from a dictionary.
 
     This just wraps a simple dictionary look-up with a helpful error that is
@@ -503,7 +503,7 @@ def _get_doc_mask(field_paths: Iterable[str]) -> Optional[types.common.DocumentM
         return types.DocumentMask(field_paths=field_paths)
 
 
-def _item_to_collection_ref(iterator, item) -> BaseCollectionReference:
+def _item_to_collection_ref(iterator, item: str) -> BaseCollectionReference:
     """Convert collection ID to collection ref.
 
     Args:
@@ -514,7 +514,7 @@ def _item_to_collection_ref(iterator, item) -> BaseCollectionReference:
     return iterator.client.collection(item)
 
 
-def _path_helper(path) -> Tuple[str]:
+def _path_helper(path: tuple) -> Tuple[str]:
     """Standardize path into a tuple of path segments.
 
     Args:

google/cloud/firestore_v1/base_collection.py

@@ -105,7 +105,7 @@ def parent(self):
     def _query(self) -> BaseQuery:
         raise NotImplementedError
 
-    def document(self, document_id=None) -> DocumentReference:
+    def document(self, document_id: str = None) -> DocumentReference:
         """Create a sub-document underneath the current collection.
 
         Args:

google/cloud/firestore_v1/base_document.py

@@ -18,11 +18,11 @@
 
 from google.cloud.firestore_v1 import _helpers
 from google.cloud.firestore_v1 import field_path as field_path_module
-from typing import Any, Dict, NoReturn, Union
 
 # Types needed only for Type Hints
 from google.cloud.firestore_v1.types import firestore
 from google.cloud.firestore_v1.types import write
+from typing import Any, Dict, Iterable, NoReturn, Union, Tuple
 
 
 class BaseDocumentReference(object):
@@ -188,12 +188,10 @@ def create(self, document_data: dict) -> NoReturn:
     def set(self, document_data: dict, merge: bool = False) -> NoReturn:
         raise NotImplementedError
 
-    def update(
-        self, field_updates: dict, option: _helpers.WriteOption = None
-    ) -> NoReturn:
+    def update(self, field_updates: dict, option=None) -> NoReturn:
         raise NotImplementedError
 
-    def delete(self, option: _helpers.WriteOption = None) -> NoReturn:
+    def delete(self, option=None) -> NoReturn:
         raise NotImplementedError
 
     def get(
@@ -431,7 +429,7 @@ def _consume_single_get(response_iterator) -> firestore.BatchGetDocumentsRespons
     return all_responses[0]
 
 
-def _first_write_result(write_results) -> write.WriteResult:
+def _first_write_result(write_results: list) -> write.WriteResult:
     """Get first write result from list.
 
     For cases where ``len(write_results) > 1``, this assumes the writes

google/cloud/firestore_v1/base_query.py

@@ -391,7 +391,7 @@ def limit(self, count: int) -> "BaseQuery":
             all_descendants=self._all_descendants,
         )
 
-    def limit_to_last(self, count) -> "BaseQuery":
+    def limit_to_last(self, count: int) -> "BaseQuery":
         """Limit a query to return the last `count` matching results.
         If the current query already has a `limit_to_last`
         set, this will override it.
@@ -856,7 +856,7 @@ def _comparator(self, doc1, doc2) -> int:
         return 0
 
 
-def _enum_from_op_string(op_string) -> int:
+def _enum_from_op_string(op_string: str) -> int:
     """Convert a string representation of a binary operator to an enum.
 
     These enums come from the protobuf message definition
@@ -899,7 +899,7 @@ def _isnan(value) -> bool:
         return False
 
 
-def _enum_from_direction(direction) -> int:
+def _enum_from_direction(direction: str) -> int:
     """Convert a string representation of a direction to an enum.
 
     Args:

google/cloud/firestore_v1/client.py

@@ -206,7 +206,10 @@ def document(self, *document_path: Tuple[str]) -> DocumentReference:
         )
 
     def get_all(
-        self, references, field_paths=None, transaction=None
+        self,
+        references: list,
+        field_paths: Iterable[str] = None,
+        transaction: Transaction = None,
     ) -> Generator[DocumentSnapshot, Any, None]:
         """Retrieve a batch of documents.
 

google/cloud/firestore_v1/document.py

@@ -26,7 +26,7 @@
 from google.cloud.firestore_v1.types import write
 from google.cloud.firestore_v1.watch import Watch
 from google.protobuf import timestamp_pb2
-from typing import Any, Generator
+from typing import Any, Callable, Generator, Iterable
 
 
 class DocumentReference(BaseDocumentReference):
@@ -78,7 +78,7 @@ def create(self, document_data) -> write.WriteResult:
         write_results = batch.commit()
         return _first_write_result(write_results)
 
-    def set(self, document_data, merge=False) -> write.WriteResult:
+    def set(self, document_data: dict, merge: bool = False) -> write.WriteResult:
         """Replace the current document in the Firestore database.
 
         A write ``option`` can be specified to indicate preconditions of
@@ -109,7 +109,7 @@ def set(self, document_data, merge=False) -> write.WriteResult:
         write_results = batch.commit()
         return _first_write_result(write_results)
 
-    def update(self, field_updates, option=None) -> write.WriteResult:
+    def update(self, field_updates: dict, option=None) -> write.WriteResult:
         """Update an existing document in the Firestore database.
 
         By default, this method verifies that the document exists on the

google/cloud/firestore_v1/transaction.py

@@ -36,11 +36,11 @@
 from google.cloud.firestore_v1 import batch
 from google.cloud.firestore_v1.document import DocumentReference
 from google.cloud.firestore_v1.query import Query
-from typing import Any, Generator, Optional
 
 # Types needed only for Type Hints
 from google.cloud.firestore_v1.base_document import DocumentSnapshot
 from google.cloud.firestore_v1.types import CommitResponse
+from typing import Any, Callable, Generator, Optional
 
 
 class Transaction(batch.WriteBatch, BaseTransaction):
@@ -140,7 +140,7 @@ def _commit(self) -> list:
         self._clean_up()
         return list(commit_response.write_results)
 
-    def get_all(self, references) -> Generator[DocumentSnapshot, Any, None]:
+    def get_all(self, references: list) -> Generator[DocumentSnapshot, Any, None]:
         """Retrieves multiple documents from Firestore.
 
         Args:
@@ -305,7 +305,9 @@ def transactional(to_wrap: Callable) -> _Transactional:
     return _Transactional(to_wrap)
 
 
-def _commit_with_retry(client, write_pbs, transaction_id) -> CommitResponse:
+def _commit_with_retry(
+    client, write_pbs: list, transaction_id: bytes
+) -> CommitResponse:
     """Call ``Commit`` on the GAPIC client with retry / sleep.
 
     Retries the ``Commit`` RPC on Unavailable. Usually this RPC-level
@@ -348,7 +350,9 @@ def _commit_with_retry(client, write_pbs, transaction_id) -> CommitResponse:
         current_sleep = _sleep(current_sleep)
 
 
-def _sleep(current_sleep, max_sleep=_MAX_SLEEP, multiplier=_MULTIPLIER) -> float:
+def _sleep(
+    current_sleep: float, max_sleep: float = _MAX_SLEEP, multiplier: float = _MULTIPLIER
+) -> float:
     """Sleep and produce a new sleep time.
 
     .. _Exponential Backoff And Jitter: https://www.awsarchitectureblog.com/\