black

Author: xiangyan99

sdk/search/azure-search-documents/azure/search/documents/aio/_search_client_async.py

@@ -66,11 +66,7 @@ class SearchClient(HeadersMixin):
     _client: SearchIndexClient
 
     def __init__(
-        self,
-        endpoint: str,
-        index_name: str,
-        credential: Union[AzureKeyCredential, AsyncTokenCredential],
-        **kwargs: Any
+        self, endpoint: str, index_name: str, credential: Union[AzureKeyCredential, AsyncTokenCredential], **kwargs: Any
     ) -> None:
         self._api_version = kwargs.pop("api_version", DEFAULT_VERSION)
         self._index_documents_batch = IndexDocumentsBatch()
@@ -89,9 +85,7 @@ def __init__(
             )
         else:
             self._aad = True
-            authentication_policy = get_authentication_policy(
-                credential, audience=audience, is_async=True
-            )
+            authentication_policy = get_authentication_policy(credential, audience=audience, is_async=True)
             self._client = SearchIndexClient(
                 endpoint=endpoint,
                 index_name=index_name,
@@ -102,9 +96,7 @@ def __init__(
             )
 
     def __repr__(self) -> str:
-        return "<SearchClient [endpoint={}, index={}]>".format(
-            repr(self._endpoint), repr(self._index_name)
-        )[:1024]
+        return "<SearchClient [endpoint={}, index={}]>".format(repr(self._endpoint), repr(self._index_name))[:1024]
 
     async def close(self) -> None:
         """Close the :class:`~azure.search.documents.aio.SearchClient` session.
@@ -125,9 +117,7 @@ async def get_document_count(self, **kwargs: Any) -> int:
         return int(await self._client.documents.count(**kwargs))
 
     @distributed_trace_async
-    async def get_document(
-        self, key: str, selected_fields: Optional[List[str]] = None, **kwargs: Any
-    ) -> Dict:
+    async def get_document(self, key: str, selected_fields: Optional[List[str]] = None, **kwargs: Any) -> Dict:
         """Retrieve a document from the Azure search index by its key.
 
         :param key: The primary key value for the document to retrieve
@@ -147,9 +137,7 @@ async def get_document(
                 :caption: Get a specific document from the search index.
         """
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
-        result = await self._client.documents.get(
-            key=key, selected_fields=selected_fields, **kwargs
-        )
+        result = await self._client.documents.get(key=key, selected_fields=selected_fields, **kwargs)
         return cast(dict, result)
 
     @distributed_trace_async
@@ -340,16 +328,8 @@ async def search(
         include_total_result_count = include_total_count
         filter_arg = filter
         search_fields_str = ",".join(search_fields) if search_fields else None
-        answers = (
-            query_answer
-            if not query_answer_count
-            else "{}|count-{}".format(query_answer, query_answer_count)
-        )
-        answers = (
-            answers
-            if not query_answer_threshold
-            else "{}|threshold-{}".format(answers, query_answer_threshold)
-        )
+        answers = query_answer if not query_answer_count else "{}|count-{}".format(query_answer, query_answer_count)
+        answers = answers if not query_answer_threshold else "{}|threshold-{}".format(answers, query_answer_threshold)
         captions = (
             query_caption
             if not query_caption_highlight
@@ -396,9 +376,7 @@ async def search(
             query.order_by(order_by)
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
         kwargs["api_version"] = self._api_version
-        return AsyncSearchItemPaged(
-            self._client, query, kwargs, page_iterator_class=AsyncSearchPageIterator
-        )
+        return AsyncSearchItemPaged(self._client, query, kwargs, page_iterator_class=AsyncSearchPageIterator)
 
     @distributed_trace_async
     async def suggest(
@@ -481,9 +459,7 @@ async def suggest(
             query.order_by(order_by)
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
         request = cast(SuggestRequest, query.request)
-        response = await self._client.documents.suggest_post(
-            suggest_request=request, **kwargs
-        )
+        response = await self._client.documents.suggest_post(suggest_request=request, **kwargs)
         assert response.results is not None  # Hint for mypy
         results = [r.as_dict() for r in response.results]
         return results
@@ -560,17 +536,13 @@ async def autocomplete(
 
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
         request = cast(AutocompleteRequest, query.request)
-        response = await self._client.documents.autocomplete_post(
-            autocomplete_request=request, **kwargs
-        )
+        response = await self._client.documents.autocomplete_post(autocomplete_request=request, **kwargs)
         assert response.results is not None  # Hint for mypy
         results = [r.as_dict() for r in response.results]
         return results
 
     # pylint:disable=client-method-missing-tracing-decorator-async
-    async def upload_documents(
-        self, documents: List[Dict], **kwargs: Any
-    ) -> List[IndexingResult]:
+    async def upload_documents(self, documents: List[Dict], **kwargs: Any) -> List[IndexingResult]:
         """Upload documents to the Azure search index.
 
         An upload action is similar to an "upsert" where the document will be
@@ -599,9 +571,7 @@ async def upload_documents(
         return cast(List[IndexingResult], results)
 
     # pylint:disable=client-method-missing-tracing-decorator-async, delete-operation-wrong-return-type
-    async def delete_documents(
-        self, documents: List[Dict], **kwargs: Any
-    ) -> List[IndexingResult]:
+    async def delete_documents(self, documents: List[Dict], **kwargs: Any) -> List[IndexingResult]:
         """Delete documents from the Azure search index
 
         Delete removes the specified document from the index. Any field you
@@ -635,9 +605,7 @@ async def delete_documents(
         return cast(List[IndexingResult], results)
 
     # pylint:disable=client-method-missing-tracing-decorator-async
-    async def merge_documents(
-        self, documents: List[Dict], **kwargs: Any
-    ) -> List[IndexingResult]:
+    async def merge_documents(self, documents: List[Dict], **kwargs: Any) -> List[IndexingResult]:
         """Merge documents in to existing documents in the Azure search index.
 
         Merge updates an existing document with the specified fields. If the
@@ -667,9 +635,7 @@ async def merge_documents(
         return cast(List[IndexingResult], results)
 
     # pylint:disable=client-method-missing-tracing-decorator-async
-    async def merge_or_upload_documents(
-        self, documents: List[Dict], **kwargs: Any
-    ) -> List[IndexingResult]:
+    async def merge_or_upload_documents(self, documents: List[Dict], **kwargs: Any) -> List[IndexingResult]:
         """Merge documents in to existing documents in the Azure search index,
         or upload them if they do not yet exist.
 
@@ -690,9 +656,7 @@ async def merge_or_upload_documents(
         return cast(List[IndexingResult], results)
 
     @distributed_trace_async
-    async def index_documents(
-        self, batch: IndexDocumentsBatch, **kwargs: Any
-    ) -> List[IndexingResult]:
+    async def index_documents(self, batch: IndexDocumentsBatch, **kwargs: Any) -> List[IndexingResult]:
         """Specify a document operations to perform as a batch.
 
         :param batch: A batch of document operations to perform.
@@ -703,17 +667,13 @@ async def index_documents(
         """
         return await self._index_documents_actions(actions=batch.actions, **kwargs)
 
-    async def _index_documents_actions(
-        self, actions: List[IndexAction], **kwargs: Any
-    ) -> List[IndexingResult]:
+    async def _index_documents_actions(self, actions: List[IndexAction], **kwargs: Any) -> List[IndexingResult]:
         error_map = {413: RequestEntityTooLargeError}
 
         kwargs["headers"] = self._merge_client_headers(kwargs.get("headers"))
         batch = IndexBatch(actions=actions)
         try:
-            batch_response = await self._client.documents.index(
-                batch=batch, error_map=error_map, **kwargs
-            )
+            batch_response = await self._client.documents.index(batch=batch, error_map=error_map, **kwargs)
             return cast(List[IndexingResult], batch_response.results)
         except RequestEntityTooLargeError:
             if len(actions) == 1: