improve type annotation

googleapis · Sep 11, 2024 · 0bd5899 · 0bd5899
1 parent 1199f79
commit 0bd5899
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 12 deletions.
diff --git a/google/cloud/firestore_v1/async_aggregation.py b/google/cloud/firestore_v1/async_aggregation.py
@@ -20,7 +20,7 @@
 """
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Optional, Union
 
 from google.api_core import gapic_v1
 from google.api_core import retry_async as retries
@@ -105,8 +105,8 @@ async def _make_stream(
         retry: Optional[retries.AsyncRetry] = gapic_v1.method.DEFAULT,
         timeout: Optional[float] = None,
         explain_options: Optional[ExplainOptions] = None,
-    ) -> AsyncStreamGenerator[
-        List[AggregationResult] | query_profile_pb.ExplainMetrics
+    ) -> AsyncGenerator[
+        [List[AggregationResult] | query_profile_pb.ExplainMetrics], Any
     ]:
         """Internal method for stream(). Runs the aggregation query.
 
@@ -134,7 +134,7 @@ async def _make_stream(
 
         Yields:
             List[AggregationResult] | query_profile_pb.ExplainMetrics:
-            The result of aggregations of this query
+            The result of aggregations of this query.
         """
         metrics: query_profile_pb.ExplainMetrics | None = None
 

diff --git a/google/cloud/firestore_v1/async_collection.py b/google/cloud/firestore_v1/async_collection.py
@@ -237,7 +237,7 @@ def stream(
         timeout: Optional[float] = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
-    ) -> "AsyncStreamGenerator[DocumentSnapshot]":
+    ) -> AsyncStreamGenerator[DocumentSnapshot]:
         """Read the documents in this collection.
 
         This sends a ``RunQuery`` RPC and then returns a generator which

diff --git a/google/cloud/firestore_v1/async_query.py b/google/cloud/firestore_v1/async_query.py
@@ -411,7 +411,7 @@ def stream(
         timeout: Optional[float] = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
-    ) -> AsyncGenerator[async_document.DocumentSnapshot, Any]:
+    ) -> AsyncStreamGenerator[async_document.DocumentSnapshot]:
         """Read the documents in the collection that match this query.
 
         This sends a ``RunQuery`` RPC and then returns a generator which
@@ -444,7 +444,7 @@ def stream(
                 explain_metrics will be available on the returned generator.
 
         Returns:
-            `AsyncGenerator[async_document.DocumentSnapshot, Any]`:
+            `AsyncStreamGenerator[async_document.DocumentSnapshot]`:
             An asynchronous generator of the queryresults.
         """
         inner_generator = self._make_stream(

diff --git a/google/cloud/firestore_v1/async_stream_generator.py b/google/cloud/firestore_v1/async_stream_generator.py
@@ -52,7 +52,7 @@ class AsyncStreamGenerator(AsyncGenerator[T, Any]):
 
     def __init__(
         self,
-        response_generator: AsyncGenerator[Any, Any],
+        response_generator: AsyncGenerator[T, Any],
         explain_options: Optional[ExplainOptions] = None,
     ):
         self._generator = response_generator

diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py
@@ -37,6 +37,7 @@
     BaseTransaction,
     _BaseTransactional,
 )
+from google.cloud.firestore_v1.query_results import QueryResultsList
 
 # Types needed only for Type Hints
 if TYPE_CHECKING:  # pragma: NO COVER
@@ -179,7 +180,7 @@ async def get(
         timeout: float = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
-    ) -> AsyncGenerator[DocumentSnapshot, Any]:
+    ) -> QueryResultsList[DocumentSnapshot]:
         """
         Retrieve a document or a query result from the database.
 
@@ -194,9 +195,9 @@ async def get(
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
 
-        Yields:
-            .DocumentSnapshot: The next document snapshot that fulfills the
-            query, or :data:`None` if the document does not exist.
+        Returns:
+            QueryResultsList[DocumentSnapshot]: The next document snapshot that
+            fulfills the query, or :data:`None` if the document does not exist.
         """
         kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
         if isinstance(ref_or_query, AsyncDocumentReference):