feat: improve type information (#176)

Co-authored-by: Tres Seaver <tseaver@palladion.com>

google/cloud/firestore_v1/_helpers.py

@@ -32,7 +32,7 @@
 from google.cloud.firestore_v1.types import common
 from google.cloud.firestore_v1.types import document
 from google.cloud.firestore_v1.types import write
-from typing import Any, Generator, List, NoReturn, Optional, Tuple
+from typing import Any, Generator, List, NoReturn, Optional, Tuple, Union
 
 _EmptyDict: transforms.Sentinel
 _GRPC_ERROR_MAPPING: dict
@@ -69,7 +69,7 @@ def __init__(self, latitude, longitude) -> None:
         self.latitude = latitude
         self.longitude = longitude
 
-    def to_protobuf(self) -> Any:
+    def to_protobuf(self) -> latlng_pb2.LatLng:
         """Convert the current object to protobuf.
 
         Returns:
@@ -253,7 +253,9 @@ def reference_value_to_document(reference_value, client) -> Any:
     return document
 
 
-def decode_value(value, client) -> Any:
+def decode_value(
+    value, client
+) -> Union[None, bool, int, float, list, datetime.datetime, str, bytes, dict, GeoPoint]:
     """Converts a Firestore protobuf ``Value`` to a native Python value.
 
     Args:
@@ -316,7 +318,7 @@ def decode_dict(value_fields, client) -> dict:
     return {key: decode_value(value, client) for key, value in value_fields.items()}
 
 
-def get_doc_id(document_pb, expected_prefix) -> Any:
+def get_doc_id(document_pb, expected_prefix) -> str:
     """Parse a document ID from a document protobuf.
 
     Args:
@@ -887,7 +889,7 @@ class ReadAfterWriteError(Exception):
     """
 
 
-def get_transaction_id(transaction, read_operation=True) -> Any:
+def get_transaction_id(transaction, read_operation=True) -> Union[bytes, None]:
     """Get the transaction ID from a ``Transaction`` object.
 
     Args:

google/cloud/firestore_v1/async_document.py

@@ -25,6 +25,8 @@
 
 from google.api_core import exceptions  # type: ignore
 from google.cloud.firestore_v1 import _helpers
+from google.cloud.firestore_v1.types import write
+from google.protobuf import timestamp_pb2
 from typing import Any, AsyncGenerator, Coroutine, Iterable, Union
 
 
@@ -61,7 +63,7 @@ async def create(
         document_data: dict,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Coroutine:
+    ) -> write.WriteResult:
         """Create the current document in the Firestore database.
 
         Args:
@@ -91,7 +93,7 @@ async def set(
         merge: bool = False,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Coroutine:
+    ) -> write.WriteResult:
         """Replace the current document in the Firestore database.
 
         A write ``option`` can be specified to indicate preconditions of
@@ -131,7 +133,7 @@ async def update(
         option: _helpers.WriteOption = None,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Coroutine:
+    ) -> write.WriteResult:
         """Update an existing document in the Firestore database.
 
         By default, this method verifies that the document exists on the
@@ -287,7 +289,7 @@ async def delete(
         option: _helpers.WriteOption = None,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Coroutine:
+    ) -> timestamp_pb2.Timestamp:
         """Delete the current document in the Firestore database.
 
         Args:

google/cloud/firestore_v1/async_transaction.py

@@ -153,7 +153,7 @@ async def get_all(
         references: list,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Coroutine:
+    ) -> AsyncGenerator[DocumentSnapshot, Any]:
         """Retrieves multiple documents from Firestore.
 
         Args:

google/cloud/firestore_v1/base_client.py

@@ -166,7 +166,7 @@ def _firestore_api_helper(self, transport, client_class, client_module) -> Any:
 
         return self._firestore_api_internal
 
-    def _target_helper(self, client_class) -> Any:
+    def _target_helper(self, client_class) -> str:
         """Return the target (where the API is).
         Eg. "firestore.googleapis.com"
 
@@ -273,7 +273,7 @@ def _document_path_helper(self, *document_path) -> List[str]:
         return joined_path.split(_helpers.DOCUMENT_PATH_DELIMITER)
 
     @staticmethod
-    def field_path(*field_names: Tuple[str]) -> Any:
+    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
@@ -438,7 +438,7 @@ def _reference_info(references: list) -> Tuple[list, dict]:
     return document_paths, reference_map
 
 
-def _get_reference(document_path: str, reference_map: dict) -> Any:
+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
@@ -536,7 +536,18 @@ def _get_doc_mask(field_paths: Iterable[str]) -> Optional[types.common.DocumentM
         return types.DocumentMask(field_paths=field_paths)
 
 
-def _path_helper(path: tuple) -> Any:
+def _item_to_collection_ref(iterator, item: str) -> BaseCollectionReference:
+    """Convert collection ID to collection ref.
+
+    Args:
+        iterator (google.api_core.page_iterator.GRPCIterator):
+            iterator response
+        item (str): ID of the collection
+    """
+    return iterator.client.collection(item)
+
+
+def _path_helper(path: tuple) -> Tuple[str]:
     """Standardize path into a tuple of path segments.
 
     Args:

google/cloud/firestore_v1/base_collection.py

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

google/cloud/firestore_v1/base_document.py

@@ -22,10 +22,10 @@
 from google.cloud.firestore_v1 import field_path as field_path_module
 from google.cloud.firestore_v1.types import common
 
-from typing import Any
-from typing import Iterable
-from typing import NoReturn
-from typing import Tuple
+# 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):
@@ -475,7 +475,7 @@ def get(self, field_path: str) -> Any:
         nested_data = field_path_module.get_nested_value(field_path, self._data)
         return copy.deepcopy(nested_data)
 
-    def to_dict(self) -> Any:
+    def to_dict(self) -> Union[Dict[str, Any], None]:
         """Retrieve the data contained in this snapshot.
 
         A copy is returned since the data may contain mutable values,
@@ -512,7 +512,7 @@ def _get_document_path(client, path: Tuple[str]) -> str:
     return _helpers.DOCUMENT_PATH_DELIMITER.join(parts)
 
 
-def _consume_single_get(response_iterator) -> Any:
+def _consume_single_get(response_iterator) -> firestore.BatchGetDocumentsResponse:
     """Consume a gRPC stream that should contain a single response.
 
     The stream will correspond to a ``BatchGetDocuments`` request made
@@ -543,7 +543,7 @@ def _consume_single_get(response_iterator) -> Any:
     return all_responses[0]
 
 
-def _first_write_result(write_results: list) -> Any:
+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

@@ -314,7 +314,7 @@ def where(self, field_path: str, op_string: str, value) -> "BaseQuery":
         )
 
     @staticmethod
-    def _make_order(field_path, direction) -> Any:
+    def _make_order(field_path, direction) -> StructuredQuery.Order:
         """Helper for :meth:`order_by`."""
         return query.StructuredQuery.Order(
             field=query.StructuredQuery.FieldReference(field_path=field_path),
@@ -394,7 +394,7 @@ def limit(self, count: int) -> "BaseQuery":
             all_descendants=self._all_descendants,
         )
 
-    def limit_to_last(self, count: int):
+    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.
@@ -651,7 +651,7 @@ def end_at(
             document_fields_or_snapshot, before=False, start=False
         )
 
-    def _filters_pb(self) -> Any:
+    def _filters_pb(self) -> StructuredQuery.Filter:
         """Convert all the filters into a single generic Filter protobuf.
 
         This may be a lone field filter or unary filter, may be a composite
@@ -674,7 +674,7 @@ def _filters_pb(self) -> Any:
             return query.StructuredQuery.Filter(composite_filter=composite_filter)
 
     @staticmethod
-    def _normalize_projection(projection) -> Any:
+    def _normalize_projection(projection) -> StructuredQuery.Projection:
         """Helper:  convert field paths to message."""
         if projection is not None:
 
@@ -836,7 +836,7 @@ def stream(
     def on_snapshot(self, callback) -> NoReturn:
         raise NotImplementedError
 
-    def _comparator(self, doc1, doc2) -> Any:
+    def _comparator(self, doc1, doc2) -> int:
         _orders = self._orders
 
         # Add implicit sorting by name, using the last specified direction.
@@ -883,7 +883,7 @@ def _comparator(self, doc1, doc2) -> Any:
         return 0
 
 
-def _enum_from_op_string(op_string: str) -> Any:
+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
@@ -926,7 +926,7 @@ def _isnan(value) -> bool:
         return False
 
 
-def _enum_from_direction(direction: str) -> Any:
+def _enum_from_direction(direction: str) -> int:
     """Convert a string representation of a direction to an enum.
 
     Args:
@@ -954,7 +954,7 @@ def _enum_from_direction(direction: str) -> Any:
         raise ValueError(msg)
 
 
-def _filter_pb(field_or_unary) -> Any:
+def _filter_pb(field_or_unary) -> StructuredQuery.Filter:
     """Convert a specific protobuf filter to the generic filter type.
 
     Args:

google/cloud/firestore_v1/client.py

@@ -46,6 +46,9 @@
 )
 from typing import Any, Generator, Iterable, Tuple
 
+# Types needed only for Type Hints
+from google.cloud.firestore_v1.base_document import DocumentSnapshot
+
 
 class Client(BaseClient):
     """Client for interacting with Google Cloud Firestore API.
@@ -209,7 +212,7 @@ def get_all(
         transaction: Transaction = None,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Generator[Any, Any, None]:
+    ) -> Generator[DocumentSnapshot, Any, None]:
         """Retrieve a batch of documents.
 
         .. note::

google/cloud/firestore_v1/document.py

@@ -25,7 +25,9 @@
 
 from google.api_core import exceptions  # type: ignore
 from google.cloud.firestore_v1 import _helpers
+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, Callable, Generator, Iterable
 
 
@@ -62,7 +64,7 @@ def create(
         document_data: dict,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Any:
+    ) -> write.WriteResult:
         """Create the current document in the Firestore database.
 
         Args:
@@ -92,7 +94,7 @@ def set(
         merge: bool = False,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Any:
+    ) -> write.WriteResult:
         """Replace the current document in the Firestore database.
 
         A write ``option`` can be specified to indicate preconditions of
@@ -132,7 +134,7 @@ def update(
         option: _helpers.WriteOption = None,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Any:
+    ) -> write.WriteResult:
         """Update an existing document in the Firestore database.
 
         By default, this method verifies that the document exists on the
@@ -288,7 +290,7 @@ def delete(
         option: _helpers.WriteOption = None,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> Any:
+    ) -> timestamp_pb2.Timestamp:
         """Delete the current document in the Firestore database.
 
         Args:
@@ -339,7 +341,7 @@ def get(
             transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
                 An existing transaction that this reference
                 will be retrieved in.
-            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+            retry (google.api_core.retry.Retry): Designation of what errors, if an                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      y,
                 should be retried.  Defaults to a system-specified policy.
             timeout (float): The timeout for this request.  Defaults to a
                 system-specified value.