feat: add _JSONSafeTypes annotation to to_dict methods

Author: Caiofcas

elasticsearch_dsl/function.py

@@ -18,7 +18,7 @@
 import collections.abc
 from typing import Dict, Optional, ClassVar
 
-from .utils import DslBase
+from .utils import DslBase, _JSONSafeTypes
 
 
 # Incomplete annotation to not break query.py tests
@@ -72,12 +72,14 @@ class ScoreFunction(DslBase):
     }
     name: ClassVar[Optional[str]] = None
 
-    def to_dict(self):
+    def to_dict(self) -> Dict[str, _JSONSafeTypes]:
         d = super().to_dict()
         # filter and query dicts should be at the same level as us
         for k in self._param_defs:
-            if k in d[self.name]:
-                d[k] = d[self.name].pop(k)
+            if self.name is not None:
+                val = d[self.name]
+                if isinstance(val, dict) and k in val:
+                    d[k] = val.pop(k)
         return d
 
 
@@ -88,12 +90,15 @@ class ScriptScore(ScoreFunction):
 class BoostFactor(ScoreFunction):
     name = "boost_factor"
 
-    def to_dict(self) -> Dict[str, int]:
+    def to_dict(self) -> Dict[str, _JSONSafeTypes]:
         d = super().to_dict()
-        if "value" in d[self.name]:
-            d[self.name] = d[self.name].pop("value")
-        else:
-            del d[self.name]
+        if self.name is not None:
+            val = d[self.name]
+            if isinstance(val, dict):
+                if "value" in val:
+                    d[self.name] = val.pop("value")
+                else:
+                    del d[self.name]
         return d
 
 

elasticsearch_dsl/utils.py

@@ -18,12 +18,14 @@
 
 import collections.abc
 from copy import copy
-from typing import Any, Dict, Optional, Type, ClassVar, Union
+from typing import Any, Dict, Optional, Type, ClassVar, Union, List
 
 from typing_extensions import Self
 
 from .exceptions import UnknownDslObject, ValidationException
 
+_JSONSafeTypes = Union[int, bool, str, float, List["_JSONSafeTypes"], Dict[str, "_JSONSafeTypes"]]
+
 SKIP_VALUES = ("", None)
 EXPAND__TO_DOT = True
 
@@ -356,8 +358,7 @@ def __getattr__(self, name):
             return AttrDict(value)
         return value
 
-    # TODO: This type annotation can probably be made tighter
-    def to_dict(self) -> Dict[str, Dict[str, Any]]:
+    def to_dict(self) -> Dict[str, _JSONSafeTypes]:
         """
         Serialize the DSL object to plain dict
         """