feat: query profiling part 1: synchronous

google/cloud/firestore_v1/async_stream_generator.py

@@ -16,25 +16,28 @@
 Firestore API.
 """
 
-from collections import abc
+from typing import Any, AsyncGenerator, Awaitable, TypeVar
 
 
-class AsyncStreamGenerator(abc.AsyncGenerator):
+T = TypeVar("T")
+
+
+class AsyncStreamGenerator(AsyncGenerator[T, Any]):
     """Asynchronous generator for the streamed results."""
 
-    def __init__(self, response_generator: abc.AsyncGenerator):
+    def __init__(self, response_generator: AsyncGenerator[T, Any]):
         self._generator = response_generator
 
-    def __aiter__(self):
-        return self._generator
+    def __aiter__(self) -> AsyncGenerator[T, Any]:
+        return self
 
-    def __anext__(self):
+    def __anext__(self) -> Awaitable[T]:
         return self._generator.__anext__()
 
-    def asend(self, value=None):
+    def asend(self, value=None) -> Awaitable[Any]:
         return self._generator.asend(value)
 
-    def athrow(self, exp=None):
+    def athrow(self, exp=None) -> Awaitable[Any]:
         return self._generator.athrow(exp)
 
     def aclose(self):

google/cloud/firestore_v1/base_aggregation.py

@@ -24,15 +24,7 @@
 
 import abc
 from abc import ABC
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Awaitable,
-    List,
-    Optional,
-    Tuple,
-    Union,
-)
+from typing import TYPE_CHECKING, List, Optional, Tuple, Union
 
 from google.api_core import gapic_v1
 from google.api_core import retry as retries
@@ -239,10 +231,7 @@ def get(
         timeout: float | None = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
-    ) -> (
-        QueryResultsList[AggregationResult]
-        | Awaitable[Any, Any, List[AggregationResult]]
-    ):
+    ) -> QueryResultsList[List[AggregationResult]] | List[List[AggregationResult]]:
         """Runs the aggregation query.
 
         This sends a ``RunAggregationQuery`` RPC and returns a list of aggregation results in the stream of ``RunAggregationQueryResponse`` messages.
@@ -264,9 +253,8 @@ def get(
                 explain_metrics will be available on the returned generator.
 
         Returns:
-            QueryResultsList[AggregationResult] | Awaitable[Any, Any, List[AggregationResult]]:
+            (QueryResultsList[List[AggregationResult]] | List[List[AggregationResult]]):
                 The aggregation query results.
-
         """
 
     @abc.abstractmethod
@@ -279,10 +267,7 @@ def stream(
         timeout: Optional[float] = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
-    ) -> (
-        StreamGenerator[List[AggregationResult]]
-        | AsyncStreamGenerator[List[AggregationResult], Any]
-    ):
+    ) -> StreamGenerator[List[AggregationResult]] | AsyncStreamGenerator:
         """Runs the aggregation query.
 
         This sends a``RunAggregationQuery`` RPC and returns a generator in the stream of ``RunAggregationQueryResponse`` messages.
@@ -302,5 +287,6 @@ def stream(
                 explain_metrics will be available on the returned generator.
 
         Returns:
+            StreamGenerator[List[AggregationResult]] | AsyncStreamGenerator:
             A generator of the query results.
         """

google/cloud/firestore_v1/base_collection.py

@@ -48,6 +48,7 @@
     from google.cloud.firestore_v1.document import DocumentReference
     from google.cloud.firestore_v1.field_path import FieldPath
     from google.cloud.firestore_v1.query_profile import ExplainOptions
+    from google.cloud.firestore_v1.query_results import QueryResultsList
     from google.cloud.firestore_v1.transaction import Transaction
     from google.cloud.firestore_v1.vector import Vector
     from google.cloud.firestore_v1.vector_query import VectorQuery
@@ -497,9 +498,7 @@ def get(
         timeout: Optional[float] = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
-    ) -> Union[
-        Generator[DocumentSnapshot, Any, Any], AsyncGenerator[DocumentSnapshot, Any]
-    ]:
+    ) -> QueryResultsList[DocumentSnapshot]:
         raise NotImplementedError
 
     def stream(

google/cloud/firestore_v1/query_results.py

@@ -18,11 +18,12 @@
     QueryExplainError,
 )
 
+from typing import List, Optional, TypeVar
 
-from typing import List, Optional
+T = TypeVar("T")
 
 
-class QueryResultsList(list):
+class QueryResultsList(List[T]):
     """A list of received query results from the query call.
 
     This is a subclass of the built-in list. A new property `explain_metrics`
@@ -68,4 +69,4 @@ def get_explain_metrics(self) -> ExplainMetrics:
         if self._explain_options is None:
             raise QueryExplainError("explain_options not set on query.")
         else:
-            return self._explain_metrics
+            return self._explain_metrics  # type: ignore

google/cloud/firestore_v1/stream_generator.py

@@ -30,7 +30,7 @@
 T = TypeVar("T")
 
 
-class StreamGenerator(Generator[T, Any, None]):
+class StreamGenerator(Generator[T, Any, Optional[ExplainMetrics]]):
     """Generator for the streamed results.
 
     Args: