address comments

google/cloud/firestore_v1/async_aggregation.py

@@ -132,7 +132,10 @@ 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. Query results will be
+            yielded as `List[AggregationResult]`. When the result contains
+            returned explain metrics, yield `query_profile_pb.ExplainMetrics`
+            individually.
         """
         metrics: query_profile_pb.ExplainMetrics | None = None
 

google/cloud/firestore_v1/async_query.py

@@ -369,7 +369,9 @@ async def _make_stream(
         Yields:
             [:class:`~google.cloud.firestore_v1.base_document.DocumentSnapshot` \
                 | google.cloud.firestore_v1.types.query_profile.ExplainMetrtics]:
-            The next document that fulfills the query.
+            The next document that fulfills the query. Query results will be
+            yielded as `DocumentSnapshot`. When the result contains returned
+            explain metrics, yield `query_profile_pb.ExplainMetrics` individually.
         """
         metrics: query_profile_pb.ExplainMetrics | None = None
 

google/cloud/firestore_v1/async_transaction.py

@@ -173,7 +173,7 @@ async def get_all(
 
     async def get(
         self,
-        ref_or_query,
+        ref_or_query: AsyncDocumentReference | AsyncQuery,
         retry: retries.AsyncRetry = gapic_v1.method.DEFAULT,
         timeout: Optional[float] = None,
         *,
@@ -183,7 +183,8 @@ async def get(
         Retrieve a document or a query result from the database.
 
         Args:
-            ref_or_query The document references or query object to return.
+            ref_or_query (AsyncDocumentReference | AsyncQuery):
+                The document references or query object to return.
             retry (google.api_core.retry.Retry): Designation of what errors, if any,
                 should be retried.  Defaults to a system-specified policy.
             timeout (float): The timeout for this request.  Defaults to a
@@ -192,6 +193,8 @@ async def get(
                 (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
+                Can only be used when running a query, not a document reference.
+
 
         Yields:
             DocumentSnapshot: The next document snapshot that fulfills the query,
@@ -211,7 +214,7 @@ async def get(
         elif isinstance(ref_or_query, AsyncQuery):
             if explain_options is not None:
                 kwargs["explain_options"] = explain_options
-            return ref_or_query.stream(transaction=self, **kwargs)
+            return await ref_or_query.stream(transaction=self, **kwargs)
         else:
             raise ValueError(
                 'Value for argument "ref_or_query" must be a AsyncDocumentReference or a AsyncQuery.'

google/cloud/firestore_v1/async_vector_query.py

@@ -142,7 +142,9 @@ async def _make_stream(
         Yields:
             [:class:`~google.cloud.firestore_v1.base_document.DocumentSnapshot` \
                 | google.cloud.firestore_v1.types.query_profile.ExplainMetrtics]:
-            The next document that fulfills the query.
+            The next document that fulfills the query. Query results will be
+            yielded as `DocumentSnapshot`. When the result contains returned
+            explain metrics, yield `query_profile_pb.ExplainMetrics` individually.
         """
         metrics: query_profile_pb.ExplainMetrics | None = None
 

google/cloud/firestore_v1/transaction.py

@@ -173,16 +173,17 @@ def get_all(
 
     def get(
         self,
-        ref_or_query,
+        ref_or_query: DocumentReference | Query,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
-        timeout: float = None,
+        timeout: Optional[float] = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
     ) -> StreamGenerator[DocumentSnapshot] | Generator[DocumentSnapshot, Any, None]:
         """Retrieve a document or a query result from the database.
 
         Args:
-            ref_or_query: The document references or query object to return.
+            ref_or_query (DocumentReference | Query):
+                The document references or query object to return.
             retry (google.api_core.retry.Retry): Designation of what errors, if any,
                 should be retried.  Defaults to a system-specified policy.
             timeout (float): The timeout for this request.  Defaults to a
@@ -191,7 +192,7 @@ def get(
                 (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
-                Can only be used when running a query.
+                Can only be used when running a query, not a document reference.
 
         Yields:
             .DocumentSnapshot: The next document snapshot that fulfills the

tests/unit/v1/test_async_transaction.py

@@ -362,7 +362,7 @@ async def _get_w_query_helper(retry=None, timeout=None, explain_options=None):
     client = AsyncMock(spec=[])
     transaction = _make_async_transaction(client)
     query = AsyncQuery(parent=AsyncMock(spec=[]))
-    query.stream = mock.Mock()
+    query.stream = AsyncMock()
     kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
 
     result = await transaction.get(