Improve typing (Azure#28216)

* Impove typing

* updates

* update

* updates

* updates

* updates

* update

* update

* Update sdk/search/azure-search-documents/azure/search/documents/_utils.py

Co-authored-by: Kashif Khan <[email protected]>

* update

* updates

* Update sdk/search/azure-search-documents/azure/search/documents/aio/_search_indexing_buffered_sender_async.py

Co-authored-by: Krista Pratico <[email protected]>

* Updates

* update

* remove _flatten_args

* add _flatten_args back

* update

* update

Co-authored-by: Kashif Khan <[email protected]>
Co-authored-by: Krista Pratico <[email protected]>

Author: 3 people

sdk/search/azure-search-documents/azure/search/documents/_headers_mixin.py

@@ -4,20 +4,15 @@
 # license information.
 # --------------------------------------------------------------------------
 
-from typing import TYPE_CHECKING
+from typing import Optional, Dict, Any
 
-if TYPE_CHECKING:
-    from typing import Optional
 
-
-class HeadersMixin(object):
+class HeadersMixin:
     @property
-    def _headers(self):
-        # type() -> dict
+    def _headers(self) -> Dict[str, Any]:
         return {"api-key": self._credential.key, "Accept": self._ODATA_ACCEPT}
 
-    def _merge_client_headers(self, headers):
-        # type(Optional[dict]) -> dict
+    def _merge_client_headers(self, headers: Optional[Dict[str, Any]]) -> Dict[str, Any]:
         if self._aad:
             return headers
         headers = headers or {}

sdk/search/azure-search-documents/azure/search/documents/_index_documents_batch.py

@@ -3,24 +3,18 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from typing import TYPE_CHECKING
+from typing import Union, List, Dict, Any
 from threading import Lock
 
 from ._generated.models import IndexAction
 
-if TYPE_CHECKING:
-    # pylint:disable=unused-import,ungrouped-imports
-    from typing import List
 
-
-def _flatten_args(args):
-    # type (Union[List[dict], List[List[dict]]]) -> List[dict]
+def _flatten_args(args: Union[List[Dict], List[List[Dict]]]) -> List[Dict]:
     if len(args) == 1 and isinstance(args[0], (list, tuple)):
         return args[0]
     return args
 
-
-class IndexDocumentsBatch(object):
+class IndexDocumentsBatch:
     """Represent a batch of update operations for documents in an Azure
     Search index.
 
@@ -29,17 +23,14 @@ class IndexDocumentsBatch(object):
 
     """
 
-    def __init__(self):
-        # type: () -> None
-        self._actions = []  # type: List[IndexAction]
+    def __init__(self) -> None:
+        self._actions: List[IndexAction] = []
         self._lock = Lock()
 
-    def __repr__(self):
-        # type: () -> str
+    def __repr__(self) -> str:
         return "<IndexDocumentsBatch [{} actions]>".format(len(self.actions))[:1024]
 
-    def add_upload_actions(self, *documents):
-        # type (Union[List[dict], List[List[dict]]]) -> List[IndexAction]
+    def add_upload_actions(self, *documents: Union[List[Dict], List[List[Dict]]]) -> List[IndexAction]:
         """Add documents to upload to the Azure search index.
 
         An upload action is similar to an "upsert" where the document will be
@@ -48,16 +39,16 @@ def add_upload_actions(self, *documents):
 
         :param documents: Documents to upload to an Azure search index. May be
          a single list of documents, or documents as individual parameters.
-        :type documents: dict or list[dict]
+        :type documents: Dict or List[Dict]
         :return: the added actions
         :rtype: List[IndexAction]
         """
         return self._extend_batch(_flatten_args(documents), "upload")
 
     def add_delete_actions(
-        self, *documents, **kwargs
-    ):  # pylint: disable=unused-argument
-        # type (Union[List[dict], List[List[dict]]]) -> List[IndexAction]
+        self, *documents: Union[List[Dict], List[List[Dict]]], **kwargs: Any
+    ) -> List[IndexAction]:
+        # pylint: disable=unused-argument
         """Add documents to delete to the Azure search index.
 
         Delete removes the specified document from the index. Any field you
@@ -71,16 +62,16 @@ def add_delete_actions(
 
         :param documents: Documents to delete from an Azure search index. May be
          a single list of documents, or documents as individual parameters.
-        :type documents: dict or list[dict]
+        :type documents: Dict or List[Dict]
         :return: the added actions
         :rtype: List[IndexAction]
         """
         return self._extend_batch(_flatten_args(documents), "delete")
 
     def add_merge_actions(
-        self, *documents, **kwargs
-    ):  # pylint: disable=unused-argument
-        # type (Union[List[dict], List[List[dict]]]) -> List[IndexAction]
+        self, *documents: Union[List[Dict], List[List[Dict]]], **kwargs: Any
+    ) -> List[IndexAction]:
+        # pylint: disable=unused-argument
         """Add documents to merge in to existing documents in the Azure search
         index.
 
@@ -91,16 +82,16 @@ def add_merge_actions(
 
         :param documents: Documents to merge into an Azure search index. May be
          a single list of documents, or documents as individual parameters.
-        :type documents: dict or list[dict]
+        :type documents: Dict or List[Dict]
         :return: the added actions
         :rtype: List[IndexAction]
         """
         return self._extend_batch(_flatten_args(documents), "merge")
 
     def add_merge_or_upload_actions(
-        self, *documents, **kwargs
-    ):  # pylint: disable=unused-argument
-        # type (Union[List[dict], List[List[dict]]]) -> List[IndexAction]
+        self, *documents: Union[List[Dict], List[List[Dict]]], **kwargs: Any
+    ) -> List[IndexAction]:
+        # pylint: disable=unused-argument
         """Add documents to merge in to existing documents in the Azure search
         index, or upload if they do not yet exist.
 
@@ -111,23 +102,21 @@ def add_merge_or_upload_actions(
         :param documents: Documents to merge or upload into an Azure search
          index. May be a single list of documents, or documents as individual
          parameters.
-        :type documents: dict or list[dict]
+        :type documents: Dict or List[Dict]
         :return: the added actions
         :rtype: List[IndexAction]
         """
         return self._extend_batch(_flatten_args(documents), "mergeOrUpload")
 
     @property
-    def actions(self):
-        # type: () -> List[IndexAction]
+    def actions(self) -> List[IndexAction]:
         """The list of currently index actions to index.
 
         :rtype: List[IndexAction]
         """
         return list(self._actions)
 
-    def dequeue_actions(self, **kwargs):  # pylint: disable=unused-argument
-        # type: () -> List[IndexAction]
+    def dequeue_actions(self, **kwargs: Any) -> List[IndexAction]:  # pylint: disable=unused-argument
         """Get the list of currently configured index actions and clear it.
 
         :rtype: List[IndexAction]
@@ -137,8 +126,10 @@ def dequeue_actions(self, **kwargs):  # pylint: disable=unused-argument
             self._actions = []
         return result
 
-    def enqueue_actions(self, new_actions, **kwargs):  # pylint: disable=unused-argument
-        # type: (Union[IndexAction, List[IndexAction]]) -> None
+    def enqueue_actions(
+            self, new_actions: Union[IndexAction, List[IndexAction]], **kwargs: Any
+    ) -> None:
+        # pylint: disable=unused-argument
         """Enqueue a list of index actions to index."""
         if isinstance(new_actions, IndexAction):
             with self._lock:
@@ -147,8 +138,7 @@ def enqueue_actions(self, new_actions, **kwargs):  # pylint: disable=unused-argu
             with self._lock:
                 self._actions.extend(new_actions)
 
-    def _extend_batch(self, documents, action_type):
-        # type: (List[dict], str) -> List[IndexAction]
+    def _extend_batch(self, documents: List[Dict], action_type: str) -> List[IndexAction]:
         new_actions = [
             IndexAction(additional_properties=document, action_type=action_type)
             for document in documents

sdk/search/azure-search-documents/azure/search/documents/_paging.py

@@ -3,19 +3,15 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from typing import TYPE_CHECKING
+from typing import List, Optional, Dict
 
 import base64
 import itertools
 import json
 
 from azure.core.paging import ItemPaged, PageIterator, ReturnType
 from ._generated.models import SearchRequest
-
-if TYPE_CHECKING:
-    # pylint:disable=unused-import,ungrouped-imports
-    from typing import Any, Union
-    from ..documents.models import AnswerResult
+from ..documents.models import AnswerResult
 
 
 def convert_search_result(result):
@@ -45,12 +41,11 @@ def unpack_continuation_token(token):
 
 
 class SearchItemPaged(ItemPaged[ReturnType]):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         super(SearchItemPaged, self).__init__(*args, **kwargs)
         self._first_page_iterator_instance = None
 
-    def __next__(self):
-        # type: () -> ReturnType
+    def __next__(self) -> ReturnType:
         if self._page_iterator is None:
             first_iterator = self._first_iterator_instance()
             self._page_iterator = itertools.chain.from_iterable(first_iterator)
@@ -61,29 +56,25 @@ def _first_iterator_instance(self):
             self._first_page_iterator_instance = self.by_page()
         return self._first_page_iterator_instance
 
-    def get_facets(self):
-        # type: () -> Union[dict, None]
+    def get_facets(self) -> Optional[Dict]:
         """Return any facet results if faceting was requested."""
         return self._first_iterator_instance().get_facets()
 
-    def get_coverage(self):
-        # type: () -> float
+    def get_coverage(self) -> float:
         """Return the coverage percentage, if `minimum_coverage` was
         specificied for the query.
 
         """
         return self._first_iterator_instance().get_coverage()
 
-    def get_count(self):
-        # type: () -> float
+    def get_count(self) -> int:
         """Return the count of results if `include_total_count` was
         set for the query.
 
         """
         return self._first_iterator_instance().get_count()
 
-    def get_answers(self):
-        # type: () -> Union[list[AnswerResult], None]
+    def get_answers(self) -> Optional[List[AnswerResult]]:
         """Return answers."""
         return self._first_iterator_instance().get_answers()
 
@@ -104,7 +95,7 @@ def wrapper(self, *args, **kw):
 
 
 class SearchPageIterator(PageIterator):
-    def __init__(self, client, initial_query, kwargs, continuation_token=None):
+    def __init__(self, client, initial_query, kwargs, continuation_token=None) -> None:
         super(SearchPageIterator, self).__init__(
             get_next=self._get_next_cb,
             extract_data=self._extract_data_cb,
@@ -128,7 +119,7 @@ def _get_next_cb(self, continuation_token):
             search_request=next_page_request, **self._kwargs
         )
 
-    def _extract_data_cb(self, response):  # pylint:disable=no-self-use
+    def _extract_data_cb(self, response):
         continuation_token = pack_continuation_token(
             response, api_version=self._api_version
         )

sdk/search/azure-search-documents/azure/search/documents/_queries.py

@@ -3,33 +3,25 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from typing import TYPE_CHECKING
-
+from typing import Type, Union
 from ._generated.models import AutocompleteRequest, SearchRequest, SuggestRequest
 
-if TYPE_CHECKING:
-    # pylint:disable=unused-import
-    from typing import Any, List, Type, Union
-
 
-class _QueryBase(object):
+class _QueryBase:
 
-    _request_type = (
+    _request_type: Union[Type[AutocompleteRequest], Type[SearchRequest], Type[SuggestRequest]] = (
         None
-    )  # type: Union[Type[AutocompleteRequest], Type[SearchRequest], Type[SuggestRequest]]
+    )
 
-    def __init__(self, **kwargs):
-        # type: (**Any) -> None
+    def __init__(self, **kwargs) -> None:
         self._request = self._request_type(**kwargs)  # pylint:disable=not-callable
 
-    def __repr__(self):
-        # type: () -> str
+    def __repr__(self) -> str:
         return "<{} [{}]>".format(self.__class__.__name__, self._request.search_text)[
             :1024
         ]
 
-    def filter(self, expression):
-        # type: (str) -> None
+    def filter(self, expression: str) -> None:
         """Add a `filter` expression for the search results.
 
         :param expression: An ODate expression of for the query filter.
@@ -58,25 +50,23 @@ class SearchQuery(_QueryBase):
 
     __doc__ = SearchRequest.__doc__
 
-    def order_by(self, *fields):
-        # type: (*str) -> None
+    def order_by(self, *fields: str) -> None:
         """Update the `orderby` property for the search results.
 
-        :param fields: An list of fields for the query result to be ordered by.
-        :type expression: str
+        :param fields: A list of fields for the query result to be ordered by.
+        :type fields: str
         :raises: ValueError
         """
         if not fields:
             raise ValueError("At least one field must be provided")
 
         self._request.order_by = ",".join(fields)
 
-    def select(self, *fields):
-        # type: (*str) -> None
+    def select(self, *fields: str) -> None:
         """Update the `select` property for the search results.
 
-        :param fields: An list of fields for the query result to return.
-        :type expression: str
+        :param fields: A list of fields for the query result to return.
+        :type fields: str
         :raises: ValueError
         """
         if not fields:
@@ -97,12 +87,11 @@ class SuggestQuery(_QueryBase):
 
     __doc__ = SuggestRequest.__doc__
 
-    def select(self, *fields):
-        # type: (*str) -> None
+    def select(self, *fields: str) -> None:
         """Update the `select` property for the search results.
 
-        :param fields: An list of fields for the query result to return.
-        :type expression: str
+        :param fields: A list of fields for the query result to return.
+        :type fields: str
         :raises: ValueError
         """
         if not fields: