refactor: add typing to query tests

Author: Caiofcas

elasticsearch_dsl/function.py

@@ -16,11 +16,13 @@
 #  under the License.
 
 import collections.abc
+from typing import Dict
 
 from .utils import DslBase
 
 
-def SF(name_or_sf, **params):
+# Incomplete annotation to not break query.py tests
+def SF(name_or_sf, **params) -> "ScoreFunction":
     # {"script_score": {"script": "_score"}, "filter": {}}
     if isinstance(name_or_sf, collections.abc.Mapping):
         if params:
@@ -86,7 +88,7 @@ class ScriptScore(ScoreFunction):
 class BoostFactor(ScoreFunction):
     name = "boost_factor"
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, int]:
         d = super().to_dict()
         if "value" in d[self.name]:
             d[self.name] = d[self.name].pop("value")

elasticsearch_dsl/query.py

@@ -28,6 +28,7 @@
     Protocol,
     TypeVar,
     Union,
+    cast,
     overload,
 )
 
@@ -58,7 +59,7 @@ def Q(name_or_query: QProxiedProtocol[_T]) -> _T: ...
 
 
 @overload
-def Q(name_or_query: str, **params: Any) -> "Query": ...
+def Q(name_or_query: str = "match_all", **params: Any) -> "Query": ...
 
 
 def Q(
@@ -92,7 +93,7 @@ def Q(
 
     # s.query = Q('filtered', query=s.query)
     if hasattr(name_or_query, "_proxied"):
-        return name_or_query._proxied
+        return cast(QProxiedProtocol[_T], name_or_query)._proxied
 
     # "match", title="python"
     return Query.get_dsl_class(name_or_query)(**params)

elasticsearch_dsl/utils.py

@@ -18,7 +18,7 @@
 
 import collections.abc
 from copy import copy
-from typing import Any, Optional, Type
+from typing import Any, Dict, Optional, Type
 
 from typing_extensions import Self
 
@@ -356,7 +356,8 @@ def __getattr__(self, name):
             return AttrDict(value)
         return value
 
-    def to_dict(self):
+    # TODO: This type annotation can probably be made tighter
+    def to_dict(self) -> Dict[str, Dict[str, Any]]:
         """
         Serialize the DSL object to plain dict
         """

mypy.ini

@@ -0,0 +1,3 @@
+[mypy-elasticsearch_dsl.query]
+# Allow reexport of SF for tests
+implicit_reexport = True

noxfile.py

@@ -29,7 +29,10 @@
     "utils/",
 )
 
-TYPED_FILES = ("elasticsearch_dsl/query.py",)
+TYPED_FILES = (
+    "elasticsearch_dsl/query.py",
+    "tests/test_query.py",
+)
 
 
 @nox.session(
@@ -85,7 +88,7 @@ def type_check(session):
     session.install("mypy", ".[develop]")
     errors = []
     popen = subprocess.Popen(
-        "mypy --strict elasticsearch_dsl",
+        "mypy --strict elasticsearch_dsl tests",
         env=session.env,
         shell=True,
         stdout=subprocess.PIPE,