type hits for tests and examples

Author: miguelgrinberg

elasticsearch_dsl/__init__.py

@@ -16,7 +16,7 @@
 #  under the License.
 
 from . import connections
-from .aggs import A
+from .aggs import A, Agg
 from .analysis import analyzer, char_filter, normalizer, token_filter, tokenizer
 from .document import AsyncDocument, Document
 from .document_base import InnerDoc, M, MetaField, mapped_field
@@ -81,7 +81,8 @@
 from .function import SF
 from .index import AsyncIndex, AsyncIndexTemplate, Index, IndexTemplate
 from .mapping import AsyncMapping, Mapping
-from .query import Q
+from .query import Q, Query
+from .response import AggResponse, Response, UpdateByQueryResponse
 from .search import (
     AsyncEmptySearch,
     AsyncMultiSearch,
@@ -99,6 +100,8 @@
 __versionstr__ = ".".join(map(str, VERSION))
 __all__ = [
     "A",
+    "Agg",
+    "AggResponse",
     "AsyncDocument",
     "AsyncEmptySearch",
     "AsyncFacetedSearch",
@@ -158,11 +161,13 @@
     "Object",
     "Percolator",
     "Q",
+    "Query",
     "Range",
     "RangeFacet",
     "RangeField",
     "RankFeature",
     "RankFeatures",
+    "Response",
     "SF",
     "ScaledFloat",
     "Search",
@@ -174,6 +179,7 @@
     "TokenCount",
     "UnknownDslObject",
     "UpdateByQuery",
+    "UpdateByQueryResponse",
     "ValidationException",
     "analyzer",
     "char_filter",

elasticsearch_dsl/_async/index.py

@@ -38,7 +38,7 @@ def __init__(
         name: str,
         template: str,
         index: Optional["AsyncIndex"] = None,
-        order: Optional[str] = None,
+        order: Optional[int] = None,
         **kwargs: Any,
     ):
         if index is None:
@@ -100,7 +100,7 @@ def as_template(
         self,
         template_name: str,
         pattern: Optional[str] = None,
-        order: Optional[str] = None,
+        order: Optional[int] = None,
     ) -> AsyncIndexTemplate:
         # TODO: should we allow pattern to be a top-level arg?
         # or maybe have an IndexPattern that allows for it and have

elasticsearch_dsl/_async/search.py

@@ -16,7 +16,16 @@
 #  under the License.
 
 import contextlib
-from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    cast,
+)
 
 from elasticsearch.exceptions import ApiError
 from elasticsearch.helpers import async_scan
@@ -68,6 +77,7 @@ async def count(self) -> int:
             query=cast(Optional[Dict[str, Any]], d.get("query", None)),
             **self._params,
         )
+
         return cast(int, resp["count"])
 
     async def execute(self, ignore_cache: bool = False) -> Response[_R]:
@@ -175,6 +185,10 @@ class AsyncMultiSearch(MultiSearchBase[_R]):
 
     _using: AsyncUsingType
 
+    if TYPE_CHECKING:
+
+        def add(self, search: AsyncSearch[_R]) -> Self: ...  # type: ignore[override]
+
     async def execute(
         self, ignore_cache: bool = False, raise_on_error: bool = True
     ) -> List[Response[_R]]:

elasticsearch_dsl/_sync/index.py

@@ -38,7 +38,7 @@ def __init__(
         name: str,
         template: str,
         index: Optional["Index"] = None,
-        order: Optional[str] = None,
+        order: Optional[int] = None,
         **kwargs: Any,
     ):
         if index is None:
@@ -94,7 +94,7 @@ def as_template(
         self,
         template_name: str,
         pattern: Optional[str] = None,
-        order: Optional[str] = None,
+        order: Optional[int] = None,
     ) -> IndexTemplate:
         # TODO: should we allow pattern to be a top-level arg?
         # or maybe have an IndexPattern that allows for it and have

elasticsearch_dsl/_sync/search.py

@@ -16,7 +16,7 @@
 #  under the License.
 
 import contextlib
-from typing import Any, Dict, Iterator, List, Optional, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, cast
 
 from elasticsearch.exceptions import ApiError
 from elasticsearch.helpers import scan
@@ -68,6 +68,7 @@ def count(self) -> int:
             query=cast(Optional[Dict[str, Any]], d.get("query", None)),
             **self._params,
         )
+
         return cast(int, resp["count"])
 
     def execute(self, ignore_cache: bool = False) -> Response[_R]:
@@ -169,6 +170,10 @@ class MultiSearch(MultiSearchBase[_R]):
 
     _using: UsingType
 
+    if TYPE_CHECKING:
+
+        def add(self, search: Search[_R]) -> Self: ...  # type: ignore[override]
+
     def execute(
         self, ignore_cache: bool = False, raise_on_error: bool = True
     ) -> List[Response[_R]]:

elasticsearch_dsl/document_base.py

@@ -166,24 +166,11 @@ def __init__(self, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]):
         field_defaults = {}
         for name in fields:
             value = None
-            if name in attrs:
-                # this field has a right-side value, which can be field
-                # instance on its own or wrapped with mapped_field()
-                value = attrs[name]
-                if isinstance(value, dict):
-                    # the mapped_field() wrapper function was used so we need
-                    # to look for the field instance and also record any
-                    # dataclass-style defaults
-                    value = attrs[name].get("_field")
-                    default_value = attrs[name].get("default") or attrs[name].get(
-                        "default_factory"
-                    )
-                    if default_value:
-                        field_defaults[name] = default_value
-            if value is None:
-                # the field does not have an explicit field instance given in
-                # a right-side assignment, so we need to figure out what field
-                # type to use from the annotation
+            required = None
+            multi = None
+            if name in annotations:
+                # the field has a type annotation, so next we try to figure out
+                # what field type we can use
                 type_ = annotations[name]
                 required = True
                 multi = False
@@ -201,24 +188,53 @@ def __init__(self, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]):
                     elif type_.__origin__ in [list, List]:
                         # List[type] -> mark instance as multi
                         multi = True
+                        required = False
                         type_ = type_.__args__[0]
                     else:
                         break
+                field = None
                 field_args: List[Any] = []
                 field_kwargs: Dict[str, Any] = {}
-                if not isinstance(type_, type):
-                    raise TypeError(f"Cannot map type {type_}")
-                elif issubclass(type_, InnerDoc):
+                if isinstance(type_, type) and issubclass(type_, InnerDoc):
                     # object or nested field
                     field = Nested if multi else Object
                     field_args = [type_]
                 elif type_ in self.type_annotation_map:
                     # use best field type for the type hint provided
                     field, field_kwargs = self.type_annotation_map[type_]
-                else:
-                    raise TypeError(f"Cannot map type {type_}")
-                field_kwargs = {"multi": multi, "required": required, **field_kwargs}
-                value = field(*field_args, **field_kwargs)
+
+                if field:
+                    field_kwargs = {
+                        "multi": multi,
+                        "required": required,
+                        **field_kwargs,
+                    }
+                    value = field(*field_args, **field_kwargs)
+
+            if name in attrs:
+                # this field has a right-side value, which can be field
+                # instance on its own or wrapped with mapped_field()
+                attr_value = attrs[name]
+                if isinstance(attr_value, dict):
+                    # the mapped_field() wrapper function was used so we need
+                    # to look for the field instance and also record any
+                    # dataclass-style defaults
+                    attr_value = attrs[name].get("_field")
+                    default_value = attrs[name].get("default") or attrs[name].get(
+                        "default_factory"
+                    )
+                    if default_value:
+                        field_defaults[name] = default_value
+                if attr_value:
+                    value = attr_value
+                    if required is not None:
+                        value._required = required
+                    if multi is not None:
+                        value._multi = multi
+
+            if value is None:
+                raise TypeError(f"Cannot map field {name}")
+
             self.mapping.field(name, value)
             if name in attrs:
                 del attrs[name]

elasticsearch_dsl/faceted_search_base.py

@@ -16,7 +16,19 @@
 #  under the License.
 
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Dict, Generic, List, Optional, Tuple, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 
 from typing_extensions import Self
 
@@ -26,10 +38,11 @@
 from .utils import _R, AttrDict
 
 if TYPE_CHECKING:
+    from .document_base import DocumentBase
     from .response.aggs import BucketData
     from .search_base import SearchBase
 
-FilterValueType = Union[str, datetime]
+FilterValueType = Union[str, datetime, Sequence[str]]
 
 __all__ = [
     "FacetedSearchBase",
@@ -51,7 +64,7 @@ class Facet(Generic[_R]):
     agg_type: str = ""
 
     def __init__(
-        self, metric: Optional[str] = None, metric_sort: str = "desc", **kwargs: Any
+        self, metric: Optional[Agg[_R]] = None, metric_sort: str = "desc", **kwargs: Any
     ):
         self.filter_values = ()
         self._params = kwargs
@@ -137,7 +150,9 @@ def add_filter(self, filter_values: List[FilterValueType]) -> Optional[Query]:
 class RangeFacet(Facet[_R]):
     agg_type = "range"
 
-    def _range_to_dict(self, range: Tuple[Any, Tuple[int, int]]) -> Dict[str, Any]:
+    def _range_to_dict(
+        self, range: Tuple[Any, Tuple[Optional[int], Optional[int]]]
+    ) -> Dict[str, Any]:
         key, _range = range
         out: Dict[str, Any] = {"key": key}
         if _range[0] is not None:
@@ -146,7 +161,11 @@ def _range_to_dict(self, range: Tuple[Any, Tuple[int, int]]) -> Dict[str, Any]:
             out["to"] = _range[1]
         return out
 
-    def __init__(self, ranges: List[Tuple[Any, Tuple[int, int]]], **kwargs: Any):
+    def __init__(
+        self,
+        ranges: Sequence[Tuple[Any, Tuple[Optional[int], Optional[int]]]],
+        **kwargs: Any,
+    ):
         super().__init__(**kwargs)
         self._params["ranges"] = list(map(self._range_to_dict, ranges))
         self._params["keyed"] = False
@@ -277,7 +296,7 @@ class FacetedResponse(Response[_R]):
         _facets: Dict[str, List[Tuple[Any, int, bool]]]
 
     @property
-    def query_string(self) -> Optional[Query]:
+    def query_string(self) -> Optional[Union[str, Query]]:
         return self._faceted_search._query
 
     @property
@@ -334,9 +353,9 @@ def search(self):
 
     """
 
-    index = None
-    doc_types = None
-    fields: List[str] = []
+    index: Optional[str] = None
+    doc_types: Optional[List[Union[str, Type["DocumentBase"]]]] = None
+    fields: Sequence[str] = []
     facets: Dict[str, Facet[_R]] = {}
     using = "default"
 
@@ -346,9 +365,9 @@ def search(self) -> "SearchBase[_R]": ...
 
     def __init__(
         self,
-        query: Optional[Query] = None,
+        query: Optional[Union[str, Query]] = None,
         filters: Dict[str, FilterValueType] = {},
-        sort: List[str] = [],
+        sort: Sequence[str] = [],
     ):
         """
         :arg query: the text to search for
@@ -383,16 +402,18 @@ def add_filter(
             ]
 
         # remember the filter values for use in FacetedResponse
-        self.filter_values[name] = filter_values
+        self.filter_values[name] = filter_values  # type: ignore[assignment]
 
         # get the filter from the facet
-        f = self.facets[name].add_filter(filter_values)
+        f = self.facets[name].add_filter(filter_values)  # type: ignore[arg-type]
         if f is None:
             return
 
         self._filters[name] = f
 
-    def query(self, search: "SearchBase[_R]", query: Query) -> "SearchBase[_R]":
+    def query(
+        self, search: "SearchBase[_R]", query: Union[str, Query]
+    ) -> "SearchBase[_R]":
         """
         Add query part to ``search``.