type hints for document, search_base and utils classes (#1851)

* type hints for document, search_base and utils classes

* Update elasticsearch_dsl/utils.py

Co-authored-by: Quentin Pradet <quentin.pradet@gmail.com>

* Update tests/test_validation.py

Co-authored-by: Quentin Pradet <quentin.pradet@gmail.com>

* feedback

---------

Co-authored-by: Quentin Pradet <quentin.pradet@gmail.com>

Author: miguelgrinberg

elasticsearch_dsl/_async/document.py

@@ -16,35 +16,45 @@
 #  under the License.
 
 import collections.abc
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
 
 from elasticsearch.exceptions import NotFoundError, RequestError
-from typing_extensions import dataclass_transform
+from typing_extensions import Self, dataclass_transform
 
 from .._async.index import AsyncIndex
 from ..async_connections import get_connection
 from ..document_base import DocumentBase, DocumentMeta, mapped_field
 from ..exceptions import IllegalOperation
-from ..utils import DOC_META_FIELDS, META_FIELDS, merge
+from ..utils import DOC_META_FIELDS, META_FIELDS, AsyncUsingType, merge
 from .search import AsyncSearch
 
+if TYPE_CHECKING:
+    from elasticsearch import AsyncElasticsearch
+
 
 class AsyncIndexMeta(DocumentMeta):
+    _index: AsyncIndex
+
     # global flag to guard us from associating an Index with the base Document
     # class, only user defined subclasses should have an _index attr
     _document_initialized = False
 
-    def __new__(cls, name, bases, attrs):
+    def __new__(
+        cls, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]
+    ) -> "AsyncIndexMeta":
         new_cls = super().__new__(cls, name, bases, attrs)
         if cls._document_initialized:
             index_opts = attrs.pop("Index", None)
             index = cls.construct_index(index_opts, bases)
             new_cls._index = index
             index.document(new_cls)
         cls._document_initialized = True
-        return new_cls
+        return cast(AsyncIndexMeta, new_cls)
 
     @classmethod
-    def construct_index(cls, opts, bases):
+    def construct_index(
+        cls, opts: Dict[str, Any], bases: Tuple[type, ...]
+    ) -> AsyncIndex:
         if opts is None:
             for b in bases:
                 if hasattr(b, "_index"):
@@ -69,12 +79,23 @@ class AsyncDocument(DocumentBase, metaclass=AsyncIndexMeta):
     Model-like class for persisting documents in elasticsearch.
     """
 
+    if TYPE_CHECKING:
+        _index: AsyncIndex
+
     @classmethod
-    def _get_connection(cls, using=None):
+    def _get_using(cls, using: Optional[AsyncUsingType] = None) -> AsyncUsingType:
+        return using or cls._index._using
+
+    @classmethod
+    def _get_connection(
+        cls, using: Optional[AsyncUsingType] = None
+    ) -> "AsyncElasticsearch":
         return get_connection(cls._get_using(using))
 
     @classmethod
-    async def init(cls, index=None, using=None):
+    async def init(
+        cls, index: Optional[str] = None, using: Optional[AsyncUsingType] = None
+    ) -> None:
         """
         Create the index and populate the mappings in elasticsearch.
         """
@@ -84,7 +105,9 @@ async def init(cls, index=None, using=None):
         await i.save(using=using)
 
     @classmethod
-    def search(cls, using=None, index=None):
+    def search(
+        cls, using: Optional[AsyncUsingType] = None, index: Optional[str] = None
+    ) -> AsyncSearch[Self]:
         """
         Create an :class:`~elasticsearch_dsl.Search` instance that will search
         over this ``Document``.
@@ -94,7 +117,13 @@ def search(cls, using=None, index=None):
         )
 
     @classmethod
-    async def get(cls, id, using=None, index=None, **kwargs):
+    async def get(
+        cls,
+        id: str,
+        using: Optional[AsyncUsingType] = None,
+        index: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Optional[Self]:
         """
         Retrieve a single document from elasticsearch using its ``id``.
 
@@ -113,7 +142,13 @@ async def get(cls, id, using=None, index=None, **kwargs):
         return cls.from_es(doc)
 
     @classmethod
-    async def exists(cls, id, using=None, index=None, **kwargs):
+    async def exists(
+        cls,
+        id: str,
+        using: Optional[AsyncUsingType] = None,
+        index: Optional[str] = None,
+        **kwargs: Any,
+    ) -> bool:
         """
         check if exists a single document from elasticsearch using its ``id``.
 
@@ -126,12 +161,18 @@ async def exists(cls, id, using=None, index=None, **kwargs):
         ``Elasticsearch.exists`` unchanged.
         """
         es = cls._get_connection(using)
-        return await es.exists(index=cls._default_index(index), id=id, **kwargs)
+        return bool(await es.exists(index=cls._default_index(index), id=id, **kwargs))
 
     @classmethod
     async def mget(
-        cls, docs, using=None, index=None, raise_on_error=True, missing="none", **kwargs
-    ):
+        cls,
+        docs: List[Dict[str, Any]],
+        using: Optional[AsyncUsingType] = None,
+        index: Optional[str] = None,
+        raise_on_error: bool = True,
+        missing: str = "none",
+        **kwargs: Any,
+    ) -> List[Optional[Self]]:
         r"""
         Retrieve multiple document by their ``id``\s. Returns a list of instances
         in the same order as requested.
@@ -160,7 +201,9 @@ async def mget(
         }
         results = await es.mget(index=cls._default_index(index), body=body, **kwargs)
 
-        objs, error_docs, missing_docs = [], [], []
+        objs: List[Optional[Self]] = []
+        error_docs: List[Self] = []
+        missing_docs: List[Self] = []
         for doc in results["docs"]:
             if doc.get("found"):
                 if error_docs or missing_docs:
@@ -186,14 +229,19 @@ async def mget(
             error_ids = [doc["_id"] for doc in error_docs]
             message = "Required routing not provided for documents %s."
             message %= ", ".join(error_ids)
-            raise RequestError(400, message, error_docs)
+            raise RequestError(400, message, error_docs)  # type: ignore
         if missing_docs:
             missing_ids = [doc["_id"] for doc in missing_docs]
             message = f"Documents {', '.join(missing_ids)} not found."
-            raise NotFoundError(404, message, {"docs": missing_docs})
+            raise NotFoundError(404, message, {"docs": missing_docs})  # type: ignore
         return objs
 
-    async def delete(self, using=None, index=None, **kwargs):
+    async def delete(
+        self,
+        using: Optional[AsyncUsingType] = None,
+        index: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
         """
         Delete the instance in elasticsearch.
 
@@ -214,23 +262,26 @@ async def delete(self, using=None, index=None, **kwargs):
             doc_meta["if_primary_term"] = self.meta["primary_term"]
 
         doc_meta.update(kwargs)
-        await es.delete(index=self._get_index(index), **doc_meta)
+        i = self._get_index(index)
+        assert i is not None
+
+        await es.delete(index=i, **doc_meta)
 
     async def update(
         self,
-        using=None,
-        index=None,
-        detect_noop=True,
-        doc_as_upsert=False,
-        refresh=False,
-        retry_on_conflict=None,
-        script=None,
-        script_id=None,
-        scripted_upsert=False,
-        upsert=None,
-        return_doc_meta=False,
-        **fields,
-    ):
+        using: Optional[AsyncUsingType] = None,
+        index: Optional[str] = None,
+        detect_noop: bool = True,
+        doc_as_upsert: bool = False,
+        refresh: bool = False,
+        retry_on_conflict: Optional[int] = None,
+        script: Optional[Union[str, Dict[str, Any]]] = None,
+        script_id: Optional[str] = None,
+        scripted_upsert: bool = False,
+        upsert: Optional[Dict[str, Any]] = None,
+        return_doc_meta: bool = False,
+        **fields: Any,
+    ) -> Any:
         """
         Partial update of the document, specify fields you wish to update and
         both the instance and the document in elasticsearch will be updated::
@@ -261,7 +312,7 @@ async def update(
 
         :return: operation result noop/updated
         """
-        body = {
+        body: Dict[str, Any] = {
             "doc_as_upsert": doc_as_upsert,
             "detect_noop": detect_noop,
         }
@@ -317,9 +368,13 @@ async def update(
             doc_meta["if_seq_no"] = self.meta["seq_no"]
             doc_meta["if_primary_term"] = self.meta["primary_term"]
 
+        i = self._get_index(index)
+        assert i is not None
+
         meta = await self._get_connection(using).update(
-            index=self._get_index(index), body=body, refresh=refresh, **doc_meta
+            index=i, body=body, refresh=refresh, **doc_meta
         )
+
         # update meta information from ES
         for k in META_FIELDS:
             if "_" + k in meta:
@@ -329,13 +384,13 @@ async def update(
 
     async def save(
         self,
-        using=None,
-        index=None,
-        validate=True,
-        skip_empty=True,
-        return_doc_meta=False,
-        **kwargs,
-    ):
+        using: Optional[AsyncUsingType] = None,
+        index: Optional[str] = None,
+        validate: bool = True,
+        skip_empty: bool = True,
+        return_doc_meta: bool = False,
+        **kwargs: Any,
+    ) -> Any:
         """
         Save the document into elasticsearch. If the document doesn't exist it
         is created, it is overwritten otherwise. Returns ``True`` if this
@@ -369,8 +424,11 @@ async def save(
             doc_meta["if_primary_term"] = self.meta["primary_term"]
 
         doc_meta.update(kwargs)
+        i = self._get_index(index)
+        assert i is not None
+
         meta = await es.index(
-            index=self._get_index(index),
+            index=i,
             body=self.to_dict(skip_empty=skip_empty),
             **doc_meta,
         )

elasticsearch_dsl/_async/index.py

@@ -15,13 +15,21 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+from typing import TYPE_CHECKING, Any, Optional
+
+from typing_extensions import Self
+
 from ..async_connections import get_connection
 from ..exceptions import IllegalOperation
 from ..index_base import IndexBase
+from ..utils import AsyncUsingType
 from .mapping import AsyncMapping
 from .search import AsyncSearch
 from .update_by_query import AsyncUpdateByQuery
 
+if TYPE_CHECKING:
+    from elastic_transport import ObjectApiResponse
+
 
 class AsyncIndexTemplate:
     def __init__(self, name, template, index=None, order=None, **kwargs):
@@ -56,7 +64,7 @@ async def save(self, using=None):
 
 
 class AsyncIndex(IndexBase):
-    def __init__(self, name, using="default"):
+    def __init__(self, name: str, using: AsyncUsingType = "default"):
         """
         :arg name: name of the index
         :arg using: connection alias to use, defaults to ``'default'``
@@ -83,7 +91,9 @@ async def load_mappings(self, using=None):
             self._name, using=using or self._using
         )
 
-    def clone(self, name=None, using=None):
+    def clone(
+        self, name: Optional[str] = None, using: Optional[AsyncUsingType] = None
+    ) -> Self:
         """
         Create a copy of the instance with another name or connection alias.
         Useful for creating multiple indices with shared configuration::
@@ -148,7 +158,7 @@ async def is_closed(self, using=None):
         )
         return state["metadata"]["indices"][self._name]["state"] == "close"
 
-    async def save(self, using=None):
+    async def save(self, using: Optional[AsyncUsingType] = None):
         """
         Sync the index definition with elasticsearch, creating the index if it
         doesn't exist and updating its settings and mappings if it does.
@@ -267,7 +277,9 @@ async def close(self, using=None, **kwargs):
             index=self._name, **kwargs
         )
 
-    async def delete(self, using=None, **kwargs):
+    async def delete(
+        self, using: Optional[AsyncUsingType] = None, **kwargs: Any
+    ) -> "ObjectApiResponse[Any]":
         """
         Deletes the index in elasticsearch.
 
@@ -278,15 +290,17 @@ async def delete(self, using=None, **kwargs):
             index=self._name, **kwargs
         )
 
-    async def exists(self, using=None, **kwargs):
+    async def exists(
+        self, using: Optional[AsyncUsingType] = None, **kwargs: Any
+    ) -> bool:
         """
         Returns ``True`` if the index already exists in elasticsearch.
 
         Any additional keyword arguments will be passed to
         ``Elasticsearch.indices.exists`` unchanged.
         """
-        return await self._get_connection(using).indices.exists(
-            index=self._name, **kwargs
+        return bool(
+            await self._get_connection(using).indices.exists(index=self._name, **kwargs)
         )
 
     async def exists_type(self, using=None, **kwargs):