refactor: add type hints to query.py + type_checking in CI

Author: Caio Fontes

elasticsearch_dsl/query.py

@@ -16,27 +16,69 @@
 #  under the License.
 
 import collections.abc
+from copy import deepcopy
 from itertools import chain
+from typing import (
+    Optional,
+    Any,
+    overload,
+    TypeVar,
+    Protocol,
+    Callable,
+    ClassVar,
+    Union,
+)
 
 # 'SF' looks unused but the test suite assumes it's available
 # from this module so others are liable to do so as well.
 from .function import SF  # noqa: F401
 from .function import ScoreFunction
 from .utils import DslBase
 
+_T = TypeVar("_T")
+_M = TypeVar("_M", bound=collections.abc.Mapping[str, Any])
 
-def Q(name_or_query="match_all", **params):
+
+class QProxiedProtocol(Protocol[_T]):
+    _proxied: _T
+
+
+@overload
+def Q(name_or_query: collections.abc.MutableMapping[str, _M]) -> "Query": ...
+
+
+@overload
+def Q(name_or_query: "Query") -> "Query": ...
+
+
+@overload
+def Q(name_or_query: QProxiedProtocol[_T]) -> _T: ...
+
+
+@overload
+def Q(name_or_query: str, **params: Any) -> "Query": ...
+
+
+def Q(
+    name_or_query: Union[
+        str,
+        "Query",
+        QProxiedProtocol[_T],
+        collections.abc.MutableMapping[str, _M],
+    ] = "match_all",
+    **params: Any
+) -> Union["Query", _T]:
     # {"match": {"title": "python"}}
-    if isinstance(name_or_query, collections.abc.Mapping):
+    if isinstance(name_or_query, collections.abc.MutableMapping):
         if params:
             raise ValueError("Q() cannot accept parameters when passing in a dict.")
         if len(name_or_query) != 1:
             raise ValueError(
                 'Q() can only accept dict with a single query ({"match": {...}}). '
                 "Instead it got (%r)" % name_or_query
             )
-        name, params = name_or_query.copy().popitem()
-        return Query.get_dsl_class(name)(_expand__to_dot=False, **params)
+        name, q_params = deepcopy(name_or_query).popitem()
+        return Query.get_dsl_class(name)(_expand__to_dot=False, **q_params)
 
     # MatchAll()
     if isinstance(name_or_query, Query):
@@ -57,26 +99,31 @@ def Q(name_or_query="match_all", **params):
 class Query(DslBase):
     _type_name = "query"
     _type_shortcut = staticmethod(Q)
-    name = None
+    name: ClassVar[Optional[str]] = None
+
+    # Add type annotations for methods not defined in every subclass
+    __ror__: ClassVar[Callable[["Query", "Query"], "Query"]]
+    __radd__: ClassVar[Callable[["Query", "Query"], "Query"]]
+    __rand__: ClassVar[Callable[["Query", "Query"], "Query"]]
 
-    def __add__(self, other):
+    def __add__(self, other: "Query") -> "Query":
         # make sure we give queries that know how to combine themselves
         # preference
         if hasattr(other, "__radd__"):
             return other.__radd__(self)
         return Bool(must=[self, other])
 
-    def __invert__(self):
+    def __invert__(self) -> "Query":
         return Bool(must_not=[self])
 
-    def __or__(self, other):
+    def __or__(self, other: "Query") -> "Query":
         # make sure we give queries that know how to combine themselves
         # preference
         if hasattr(other, "__ror__"):
             return other.__ror__(self)
         return Bool(should=[self, other])
 
-    def __and__(self, other):
+    def __and__(self, other: "Query") -> "Query":
         # make sure we give queries that know how to combine themselves
         # preference
         if hasattr(other, "__rand__"):
@@ -87,17 +134,17 @@ def __and__(self, other):
 class MatchAll(Query):
     name = "match_all"
 
-    def __add__(self, other):
+    def __add__(self, other: "Query") -> "Query":
         return other._clone()
 
     __and__ = __rand__ = __radd__ = __add__
 
-    def __or__(self, other):
+    def __or__(self, other: "Query") -> "MatchAll":
         return self
 
     __ror__ = __or__
 
-    def __invert__(self):
+    def __invert__(self) -> "MatchNone":
         return MatchNone()
 
 
@@ -107,17 +154,17 @@ def __invert__(self):
 class MatchNone(Query):
     name = "match_none"
 
-    def __add__(self, other):
+    def __add__(self, other: "Query") -> "MatchNone":
         return self
 
     __and__ = __rand__ = __radd__ = __add__
 
-    def __or__(self, other):
+    def __or__(self, other: "Query") -> "Query":
         return other._clone()
 
     __ror__ = __or__
 
-    def __invert__(self):
+    def __invert__(self) -> MatchAll:
         return MatchAll()
 
 
@@ -130,7 +177,7 @@ class Bool(Query):
         "filter": {"type": "query", "multi": True},
     }
 
-    def __add__(self, other):
+    def __add__(self, other: Query) -> "Bool":
         q = self._clone()
         if isinstance(other, Bool):
             q.must += other.must
@@ -143,7 +190,7 @@ def __add__(self, other):
 
     __radd__ = __add__
 
-    def __or__(self, other):
+    def __or__(self, other: Query) -> Query:
         for q in (self, other):
             if isinstance(q, Bool) and not any(
                 (q.must, q.must_not, q.filter, getattr(q, "minimum_should_match", None))
@@ -168,20 +215,20 @@ def __or__(self, other):
     __ror__ = __or__
 
     @property
-    def _min_should_match(self):
+    def _min_should_match(self) -> int:
         return getattr(
             self,
             "minimum_should_match",
             0 if not self.should or (self.must or self.filter) else 1,
         )
 
-    def __invert__(self):
+    def __invert__(self) -> Query:
         # Because an empty Bool query is treated like
         # MatchAll the inverse should be MatchNone
         if not any(chain(self.must, self.filter, self.should, self.must_not)):
             return MatchNone()
 
-        negations = []
+        negations: list[Query] = []
         for q in chain(self.must, self.filter):
             negations.append(~q)
 
@@ -195,7 +242,7 @@ def __invert__(self):
             return negations[0]
         return Bool(should=negations)
 
-    def __and__(self, other):
+    def __and__(self, other: Query) -> Query:
         q = self._clone()
         if isinstance(other, Bool):
             q.must += other.must
@@ -247,7 +294,7 @@ class FunctionScore(Query):
         "functions": {"type": "score_function", "multi": True},
     }
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any):
         if "functions" in kwargs:
             pass
         else:

elasticsearch_dsl/utils.py

@@ -18,6 +18,7 @@
 
 import collections.abc
 from copy import copy
+from typing import Optional, Any, Self
 
 from .exceptions import UnknownDslObject, ValidationException
 
@@ -251,7 +252,7 @@ class DslBase(metaclass=DslMeta):
     _param_defs = {}
 
     @classmethod
-    def get_dsl_class(cls, name, default=None):
+    def get_dsl_class(cls: type[Self], name: str, default:Optional[str]=None) -> type[Self]:
         try:
             return cls._classes[name]
         except KeyError:
@@ -261,7 +262,7 @@ def get_dsl_class(cls, name, default=None):
                 f"DSL class `{name}` does not exist in {cls._type_name}."
             )
 
-    def __init__(self, _expand__to_dot=None, **params):
+    def __init__(self, _expand__to_dot: Optional[bool]=None, **params: Any) -> None:
         if _expand__to_dot is None:
             _expand__to_dot = EXPAND__TO_DOT
         self._params = {}
@@ -390,7 +391,7 @@ def to_dict(self):
             d[pname] = value
         return {self.name: d}
 
-    def _clone(self):
+    def _clone(self) -> Self:
         c = self.__class__()
         for attr in self._params:
             c._params[attr] = copy(self._params[attr])

noxfile.py

@@ -16,6 +16,8 @@
 #  under the License.
 
 import nox
+import subprocess
+
 
 SOURCE_FILES = (
     "setup.py",
@@ -27,6 +29,10 @@
     "utils/",
 )
 
+TYPED_FILES = (
+    "elasticsearch_dsl/query.py",
+)
+
 
 @nox.session(
     python=[
@@ -76,6 +82,31 @@ def lint(session):
     session.run("python", "utils/license-headers.py", "check", *SOURCE_FILES)
 
 
+@nox.session(python="3.12")
+def type_check(session):
+    session.install("mypy", ".[develop]")
+    errors = []
+    popen = subprocess.Popen(
+        "mypy --strict elasticsearch_dsl",
+        env=session.env,
+        shell=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+
+    mypy_output = ""
+    while popen.poll() is None:
+        mypy_output += popen.stdout.read(8192).decode()
+    mypy_output += popen.stdout.read().decode()
+
+    for line in mypy_output.split("\n"):
+        filepath = line.partition(":")[0]
+        if filepath in TYPED_FILES:
+            errors.append(line)
+    if errors:
+        session.error("\n" + "\n".join(sorted(set(errors))))
+
+
 @nox.session()
 def docs(session):
     session.install(".[develop]")