typing of the update by query response

Author: miguelgrinberg

elasticsearch_dsl/response/__init__.py

@@ -215,8 +215,45 @@ def __iter__(self) -> Iterator["Agg"]:  # type: ignore[override]
 
 
 class UpdateByQueryResponse(AttrDict[Any], Generic[_R]):
+    """An Elasticsearch response.
+
+    :arg batches:
+    :arg failures:
+    :arg noops:
+    :arg deleted:
+    :arg requests_per_second:
+    :arg retries:
+    :arg task:
+    :arg timed_out:
+    :arg took:
+    :arg total:
+    :arg updated:
+    :arg version_conflicts:
+    :arg throttled:
+    :arg throttled_millis:
+    :arg throttled_until:
+    :arg throttled_until_millis:
+    """
+
     _search: "UpdateByQueryBase[_R]"
 
+    batches: int
+    failures: Sequence["types.BulkIndexByScrollFailure"]
+    noops: int
+    deleted: int
+    requests_per_second: float
+    retries: "types.Retries"
+    task: Union[str, int]
+    timed_out: bool
+    took: Any
+    total: int
+    updated: int
+    version_conflicts: int
+    throttled: Any
+    throttled_millis: Any
+    throttled_until: Any
+    throttled_until_millis: Any
+
     def __init__(
         self,
         search: "Request[_R]",

elasticsearch_dsl/types.py

@@ -4458,6 +4458,22 @@ class AggregationProfileDelegateDebugFilter(AttrDict[Any]):
     segments_counted_in_constant_time: int
 
 
+class BulkIndexByScrollFailure(AttrDict[Any]):
+    """
+    :arg cause: (required)
+    :arg id: (required)
+    :arg index: (required)
+    :arg status: (required)
+    :arg type: (required)
+    """
+
+    cause: "ErrorCause"
+    id: str
+    index: str
+    status: int
+    type: str
+
+
 class ClusterDetails(AttrDict[Any]):
     """
     :arg status: (required)
@@ -4884,6 +4900,16 @@ class QueryProfile(AttrDict[Any]):
     children: Sequence["QueryProfile"]
 
 
+class Retries(AttrDict[Any]):
+    """
+    :arg bulk: (required)
+    :arg search: (required)
+    """
+
+    bulk: int
+    search: int
+
+
 class SearchProfile(AttrDict[Any]):
     """
     :arg collector: (required)

utils/generator.py

@@ -638,8 +638,12 @@ def interface_to_python_class(
         ```
         """
         type_ = self.find_type(interface, namespace)
-        if type_["kind"] != "interface":
+        if type_["kind"] not in ["interface", "response"]:
             raise RuntimeError(f"Type {interface} is not an interface")
+        if type_["kind"] == "response":
+            # we consider responses as interfaces because they also have properties
+            # but the location of the properties is different
+            type_ = type_["body"]
         k = {"name": interface, "for_response": for_response, "args": []}
         self.add_behaviors(
             type_, k, for_types_py=for_types_py, for_response=for_response
@@ -718,14 +722,22 @@ def generate_response_init_py(schema, filename):
     """Generate response/__init__.py with all the response properties
     documented and typed.
     """
-    response_interface = schema.interface_to_python_class(
+    search_response = schema.interface_to_python_class(
         "ResponseBody",
         "_global.search",
         for_types_py=False,
         for_response=True,
     )
+    ubq_response = schema.interface_to_python_class(
+        "Response",
+        "_global.update_by_query",
+        for_types_py=False,
+        for_response=True,
+    )
     with open(filename, "wt") as f:
-        f.write(response_init_py.render(response=response_interface))
+        f.write(
+            response_init_py.render(response=search_response, ubq_response=ubq_response)
+        )
     print(f"Generated {filename}.")
 
 

utils/templates/response.__init__.py.tpl

@@ -193,8 +193,20 @@ class AggResponse(AttrDict[Any], Generic[_R]):
 
 
 class UpdateByQueryResponse(AttrDict[Any], Generic[_R]):
+    """An Elasticsearch response.
+
+    {% for arg in ubq_response.args %}
+        {% for line in arg.doc %}
+    {{ line }}
+        {% endfor %}
+    {% endfor %}
+    """
     _search: "UpdateByQueryBase[_R]"
 
+    {% for arg in ubq_response.args %}
+    {{ arg.name }}: {{ arg.type }}
+    {% endfor %}
+
     def __init__(
         self,
         search: "Request[_R]",