utils/generator.py

#  Licensed to Elasticsearch B.V. under one or more contributor
#  license agreements. See the NOTICE file distributed with
#  this work for additional information regarding copyright
#  ownership. Elasticsearch B.V. licenses this file to you under
#  the Apache License, Version 2.0 (the "License"); you may
#  not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
# 	http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing,
#  software distributed under the License is distributed on an
#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#  KIND, either express or implied.  See the License for the
#  specific language governing permissions and limitations
#  under the License.

import json
import re
import textwrap
from urllib.error import HTTPError
from urllib.request import urlopen

from jinja2 import Environment, PackageLoader, select_autoescape

from elasticsearch_dsl import VERSION

jinja_env = Environment(
    loader=PackageLoader("utils"),
    autoescape=select_autoescape(),
    trim_blocks=True,
    lstrip_blocks=True,
)
query_py = jinja_env.get_template("query.py.tpl")
aggs_py = jinja_env.get_template("aggs.py.tpl")
response_init_py = jinja_env.get_template("response.__init__.py.tpl")
types_py = jinja_env.get_template("types.py.tpl")

# map with name replacements for Elasticsearch attributes
PROP_REPLACEMENTS = {"from": "from_"}

# map with Elasticsearch type replacements
# keys and values are in given in "{namespace}:{name}" format
TYPE_REPLACEMENTS = {
    "_types.query_dsl:DistanceFeatureQuery": "_types.query_dsl:DistanceFeatureQueryBase",
}

# some aggregation types are complicated to determine from the schema, so they
# have their correct type here
AGG_TYPES = {
    "bucket_count_ks_test": "Pipeline",
    "bucket_correlation": "Pipeline",
    "bucket_sort": "Bucket",
    "categorize_text": "Bucket",
    "filter": "Bucket",
    "moving_avg": "Pipeline",
    "variable_width_histogram": "Bucket",
}


def property_to_class_name(name):
    return "".join([w.title() if w != "ip" else "IP" for w in name.split("_")])


def wrapped_doc(text, width=70, initial_indent="", subsequent_indent=""):
    """Formats a docstring as a list of lines of up to the request width."""
    return textwrap.wrap(
        text.replace("\n", " "),
        width=width,
        initial_indent=initial_indent,
        subsequent_indent=subsequent_indent,
    )


def add_dict_type(type_):
    """Add Dict[str, Any] to a Python type hint."""
    if type_.startswith("Union["):
        type_ = f"{type_[:-1]}, Dict[str, Any]]"
    else:
        type_ = f"Union[{type_}, Dict[str, Any]]"
    return type_


def add_seq_dict_type(type_):
    """Add Sequence[Dict[str, Any]] to a Python type hint."""
    if type_.startswith("Union["):
        type_ = f"{type_[:-1]}, Sequence[Dict[str, Any]]]"
    else:
        type_ = f"Union[{type_}, Sequence[Dict[str, Any]]]"
    return type_


def add_not_set(type_):
    """Add DefaultType to a Python type hint."""
    if type_.startswith("Union["):
        type_ = f'{type_[:-1]}, "DefaultType"]'
    else:
        type_ = f'Union[{type_}, "DefaultType"]'
    return type_


def type_for_types_py(type_):
    """Converts a type rendered in a generic way to the format needed in the
    types.py module.
    """
    type_ = type_.replace('"DefaultType"', "DefaultType")
    type_ = type_.replace('"InstrumentedField"', "InstrumentedField")
    type_ = re.sub(r'"(function\.[a-zA-Z0-9_]+)"', r"\1", type_)
    type_ = re.sub(r'"types\.([a-zA-Z0-9_]+)"', r'"\1"', type_)
    type_ = re.sub(r'"(wrappers\.[a-zA-Z0-9_]+)"', r"\1", type_)
    return type_


class ElasticsearchSchema:
    """Operations related to the Elasticsearch schema."""

    def __init__(self):
        response = None
        for branch in [f"{VERSION[0]}.{VERSION[1]}", "main"]:
            url = f"https://raw.githubusercontent.com/elastic/elasticsearch-specification/{branch}/output/schema/schema.json"
            try:
                response = urlopen(url)
                print(f"Initializing code generation with '{branch}' specification.")
                break
            except HTTPError:
                continue
        if not response:
            raise RuntimeError("Could not download Elasticsearch schema")
        self.schema = json.loads(response.read())

        # Interfaces collects interfaces that are seen while traversing the schema.
        # Any interfaces collected here are then rendered as Python in the
        # types.py module.
        self.interfaces = []
        self.response_interfaces = []

    def find_type(self, name, namespace=None):
        for t in self.schema["types"]:
            if t["name"]["name"] == name and (
                namespace is None or t["name"]["namespace"] == namespace
            ):
                return t

    def inherits_from(self, type_, name, namespace=None):
        while "inherits" in type_:
            type_ = self.find_type(
                type_["inherits"]["type"]["name"],
                type_["inherits"]["type"]["namespace"],
            )
            if type_["name"]["name"] == name and (
                namespace is None or type_["name"]["namespace"] == namespace
            ):
                return True
        return False

    def get_python_type(self, schema_type, for_response=False):
        """Obtain Python typing details for a given schema type

        This method returns a tuple. The first element is a string with the
        Python type hint. The second element is a dictionary with Python DSL
        specific typing details to be stored in the DslBase._param_defs
        attribute (or None if the type does not need to be in _param_defs).

        When `for_response` is `False`, any new interfaces that are discovered
        are registered to be generated in "request" style, with alternative
        Dict type hints and default values. If `for_response` is `True`,
        interfaces are generated just with their declared type, without
        Dict alternative and without defaults, to help type checkers be more
        effective at parsing response expressions.
        """
        if schema_type["kind"] == "instance_of":
            type_name = schema_type["type"]
            if type_name["namespace"] in ["_types", "internal", "_builtins"]:
                if type_name["name"] in ["integer", "uint", "long", "ulong"]:
                    return "int", None
                elif type_name["name"] in ["number", "float", "double"]:
                    return "float", None
                elif type_name["name"] == "string":
                    return "str", None
                elif type_name["name"] == "boolean":
                    return "bool", None
                elif type_name["name"] == "binary":
                    return "bytes", None
                elif type_name["name"] == "null":
                    return "None", None
                elif type_name["name"] == "Field":
                    if for_response:
                        return "str", None
                    else:
                        return 'Union[str, "InstrumentedField"]', None
                else:
                    # not an instance of a native type, so we get the type and try again
                    return self.get_python_type(
                        self.find_type(type_name["name"], type_name["namespace"]),
                        for_response=for_response,
                    )
            elif (
                type_name["namespace"] == "_types.query_dsl"
                and type_name["name"] == "QueryContainer"
            ):
                # QueryContainer maps to the DSL's Query class
                return "Query", {"type": "query"}
            elif (
                type_name["namespace"] == "_types.query_dsl"
                and type_name["name"] == "FunctionScoreContainer"
            ):
                # FunctionScoreContainer maps to the DSL's ScoreFunction class
                return "ScoreFunction", {"type": "score_function"}
            elif (
                type_name["namespace"] == "_types.aggregations"
                and type_name["name"] == "Buckets"
            ):
                if for_response:
                    return "Union[Sequence[Any], Dict[str, Any]]", None
                else:
                    return "Dict[str, Query]", {"type": "query", "hash": True}
            elif (
                type_name["namespace"] == "_types.aggregations"
                and type_name["name"] == "CompositeAggregationSource"
            ):
                # QueryContainer maps to the DSL's Query class
                return "Agg[_R]", None
            else:
                # for any other instances we get the type and recurse
                type_ = self.find_type(type_name["name"], type_name["namespace"])
                if type_:
                    return self.get_python_type(type_, for_response=for_response)

        elif schema_type["kind"] == "type_alias":
            # for an alias, we use the aliased type
            return self.get_python_type(schema_type["type"], for_response=for_response)

        elif schema_type["kind"] == "array_of":
            # for arrays we use Sequence[element_type]
            type_, param = self.get_python_type(
                schema_type["value"], for_response=for_response
            )
            return f"Sequence[{type_}]", {**param, "multi": True} if param else None

        elif schema_type["kind"] == "dictionary_of":
            # for dicts we use Mapping[key_type, value_type]
            key_type, key_param = self.get_python_type(
                schema_type["key"], for_response=for_response
            )
            value_type, value_param = self.get_python_type(
                schema_type["value"], for_response=for_response
            )
            return f"Mapping[{key_type}, {value_type}]", (
                {**value_param, "hash": True} if value_param else None
            )

        elif schema_type["kind"] == "union_of":
            if (
                len(schema_type["items"]) == 2
                and schema_type["items"][0]["kind"] == "instance_of"
                and schema_type["items"][1]["kind"] == "array_of"
                and schema_type["items"][0] == schema_type["items"][1]["value"]
            ):
                # special kind of unions in the form Union[type, Sequence[type]]
                type_, param = self.get_python_type(
                    schema_type["items"][0], for_response=for_response
                )
                if schema_type["items"][0]["type"]["name"] in [
                    "CompletionSuggestOption",
                    "PhraseSuggestOption",
                    "TermSuggestOption",
                ]:
                    # for suggest types we simplify this type and return just the array form
                    return (
                        f"Sequence[{type_}]",
                        ({"type": param["type"], "multi": True} if param else None),
                    )
                else:
                    # for every other types we produce an union with the two alternatives
                    return (
                        f"Union[{type_}, Sequence[{type_}]]",
                        ({"type": param["type"], "multi": True} if param else None),
                    )
            elif (
                len(schema_type["items"]) == 2
                and schema_type["items"][0]["kind"] == "instance_of"
                and schema_type["items"][1]["kind"] == "instance_of"
                and schema_type["items"][0]["type"]
                == {"name": "T", "namespace": "_spec_utils.PipeSeparatedFlags"}
                and schema_type["items"][1]["type"]
                == {"name": "string", "namespace": "_builtins"}
            ):
                # for now we treat PipeSeparatedFlags as a special case
                if "PipeSeparatedFlags" not in self.interfaces:
                    self.interfaces.append("PipeSeparatedFlags")
                return '"types.PipeSeparatedFlags"', None
            else:
                # generic union type
                types = list(
                    dict.fromkeys(  # eliminate duplicates
                        [
                            self.get_python_type(t, for_response=for_response)
                            for t in schema_type["items"]
                        ]
                    )
                )
                return "Union[" + ", ".join([type_ for type_, _ in types]) + "]", None

        elif schema_type["kind"] == "enum":
            # enums are mapped to Literal[member, ...]
            return (
                "Literal["
                + ", ".join(
                    [f"\"{member['name']}\"" for member in schema_type["members"]]
                )
                + "]",
                None,
            )

        elif schema_type["kind"] == "interface":
            if schema_type["name"]["namespace"] == "_types.query_dsl":
                # handle specific DSL classes explicitly to map to existing
                # Python DSL classes
                if schema_type["name"]["name"].endswith("RangeQuery"):
                    return '"wrappers.Range[Any]"', None
                elif schema_type["name"]["name"].endswith("ScoreFunction"):
                    # When dropping Python 3.8, use `removesuffix("Function")` instead
                    name = schema_type["name"]["name"][:-8]
                    return f'"function.{name}"', None
                elif schema_type["name"]["name"].endswith("DecayFunction"):
                    return '"function.DecayFunction"', None
                elif schema_type["name"]["name"].endswith("Function"):
                    return f"\"function.{schema_type['name']['name']}\"", None
            elif schema_type["name"]["namespace"] == "_types.analysis" and schema_type[
                "name"
            ]["name"].endswith("Analyzer"):
                # not expanding analyzers at this time, maybe in the future
                return "str, Dict[str, Any]", None

            # to handle other interfaces we generate a type of the same name
            # and add the interface to the interfaces.py module
            if schema_type["name"]["name"] not in self.interfaces:
                self.interfaces.append(schema_type["name"]["name"])
                if for_response:
                    self.response_interfaces.append(schema_type["name"]["name"])
            return f"\"types.{schema_type['name']['name']}\"", None
        elif schema_type["kind"] == "user_defined_value":
            # user_defined_value maps to Python's Any type
            return "Any", None

        raise RuntimeError(f"Cannot find Python type for {schema_type}")

    def add_attribute(self, k, arg, for_types_py=False, for_response=False):
        """Add an attribute to the internal representation of a class.

        This method adds the argument `arg` to the data structure for a class
        stored in `k`. In particular, the argument is added to the `k["args"]`
        list, making sure required arguments are first in the list. If the
        argument is of a type that needs Python DSL specific typing details to
        be stored in the DslBase._param_defs attribute, then this is added to
        `k["params"]`.

        When `for_types_py` is `True`, type hints are formatted in the most
        convenient way for the types.py file. When possible, double quotes are
        removed from types, and for types that are in the same file the quotes
        are kept to prevent forward references, but the "types." namespace is
        removed. When `for_types_py` is `False`, all non-native types use
        quotes and are namespaced.

        When `for_response` is `True`, type hints are not given the optional
        dictionary representation, nor the `DefaultType` used for omitted
        attributes.
        """
        try:
            type_, param = self.get_python_type(arg["type"], for_response=for_response)
        except RuntimeError:
            type_ = "Any"
            param = None
        if not for_response:
            if type_ != "Any":
                if 'Sequence["types.' in type_:
                    type_ = add_seq_dict_type(type_)  # interfaces can be given as dicts
                elif "types." in type_:
                    type_ = add_dict_type(type_)  # interfaces can be given as dicts
                type_ = add_not_set(type_)
        if for_types_py:
            type_ = type_for_types_py(type_)
        required = "(required) " if arg["required"] else ""
        server_default = (
            f" Defaults to `{arg['serverDefault']}` if omitted."
            if arg.get("serverDefault")
            else ""
        )
        doc = wrapped_doc(
            f":arg {arg['name']}: {required}{arg.get('description', '')}{server_default}",
            subsequent_indent="    ",
        )
        arg = {
            "name": PROP_REPLACEMENTS.get(arg["name"], arg["name"]),
            "type": type_,
            "doc": doc,
            "required": arg["required"],
        }
        if param is not None:
            param = {"name": arg["name"], "param": param}
        if arg["required"]:
            # insert in the right place so that all required arguments
            # appear at the top of the argument list
            i = 0
            for i in range(len(k["args"]) + 1):
                if i == len(k["args"]):
                    break
                if k["args"][i].get("positional"):
                    continue
                if k["args"][i]["required"] is False:
                    break
            k["args"].insert(i, arg)
        else:
            k["args"].append(arg)
        if param and "params" in k:
            k["params"].append(param)

    def add_behaviors(self, type_, k, for_types_py=False, for_response=False):
        """Add behaviors reported in the specification of the given type to the
        class representation.
        """
        if "behaviors" in type_:
            for behavior in type_["behaviors"]:
                if (
                    behavior["type"]["name"] != "AdditionalProperty"
                    or behavior["type"]["namespace"] != "_spec_utils"
                ):
                    # we do not support this behavior, so we ignore it
                    continue
                key_type, _ = self.get_python_type(
                    behavior["generics"][0], for_response=for_response
                )
                if "InstrumentedField" in key_type:
                    value_type, _ = self.get_python_type(
                        behavior["generics"][1], for_response=for_response
                    )
                    if for_types_py:
                        value_type = value_type.replace('"DefaultType"', "DefaultType")
                        value_type = value_type.replace(
                            '"InstrumentedField"', "InstrumentedField"
                        )
                        value_type = re.sub(
                            r'"(function\.[a-zA-Z0-9_]+)"', r"\1", value_type
                        )
                        value_type = re.sub(
                            r'"types\.([a-zA-Z0-9_]+)"', r'"\1"', value_type
                        )
                        value_type = re.sub(
                            r'"(wrappers\.[a-zA-Z0-9_]+)"', r"\1", value_type
                        )
                    k["args"].append(
                        {
                            "name": "_field",
                            "type": add_not_set(key_type),
                            "doc": [":arg _field: The field to use in this query."],
                            "required": False,
                            "positional": True,
                        }
                    )
                    k["args"].append(
                        {
                            "name": "_value",
                            "type": add_not_set(add_dict_type(value_type)),
                            "doc": [":arg _value: The query value for the field."],
                            "required": False,
                            "positional": True,
                        }
                    )
                    k["is_single_field"] = True
                else:
                    raise RuntimeError(
                        f"Non-field AdditionalProperty are not supported for interface {type_['name']['namespace']}:{type_['name']['name']}."
                    )

    def property_to_python_class(self, p):
        """Return a dictionary with template data necessary to render a schema
        property as a Python class.

        Used for "container" sub-classes such as `QueryContainer`, where each
        sub-class is represented by a Python DSL class.

        The format is as follows:

        ```python
        {
            "property_name": "the name of the property",
            "name": "the class name to use for the property",
            "docstring": "the formatted docstring as a list of strings",
            "args": [  # a Python description of each class attribute
                "name": "the name of the attribute",
                "type": "the Python type hint for the attribute",
                "doc": ["formatted lines of documentation to add to class docstring"],
                "required": bool,
                "positional": bool,
            ],
            "params": [
                "name": "the attribute name",
                "param": "the param dictionary to include in `_param_defs` for the class",
            ],  # a DSL-specific description of interesting attributes
            "is_single_field": bool  # True for single-key dicts with field key
            "is_multi_field": bool  # True for multi-key dicts with field keys
        }
        ```
        """
        k = {
            "property_name": p["name"],
            "name": property_to_class_name(p["name"]),
        }
        k["docstring"] = wrapped_doc(p.get("description") or "")
        other_classes = []
        kind = p["type"]["kind"]
        if kind == "instance_of":
            namespace = p["type"]["type"]["namespace"]
            name = p["type"]["type"]["name"]
            if f"{namespace}:{name}" in TYPE_REPLACEMENTS:
                namespace, name = TYPE_REPLACEMENTS[f"{namespace}:{name}"].split(":")
            if name == "QueryContainer" and namespace == "_types.query_dsl":
                type_ = {
                    "kind": "interface",
                    "properties": [p],
                }
            else:
                type_ = self.find_type(name, namespace)
            if p["name"] in AGG_TYPES:
                k["parent"] = AGG_TYPES[p["name"]]

            if type_["kind"] == "interface":
                # set the correct parent for bucket and pipeline aggregations
                if self.inherits_from(
                    type_, "PipelineAggregationBase", "_types.aggregations"
                ):
                    k["parent"] = "Pipeline"
                elif self.inherits_from(
                    type_, "BucketAggregationBase", "_types.aggregations"
                ):
                    k["parent"] = "Bucket"

                # generate class attributes
                k["args"] = []
                k["params"] = []
                self.add_behaviors(type_, k)
                while True:
                    for arg in type_["properties"]:
                        self.add_attribute(k, arg)
                    if "inherits" in type_ and "type" in type_["inherits"]:
                        type_ = self.find_type(
                            type_["inherits"]["type"]["name"],
                            type_["inherits"]["type"]["namespace"],
                        )
                    else:
                        break

            elif type_["kind"] == "type_alias":
                if type_["type"]["kind"] == "union_of":
                    # for unions we create sub-classes
                    for other in type_["type"]["items"]:
                        other_class = self.interface_to_python_class(
                            other["type"]["name"],
                            other["type"]["namespace"],
                            for_types_py=False,
                        )
                        other_class["parent"] = k["name"]
                        other_classes.append(other_class)
                else:
                    raise RuntimeError(
                        "Cannot generate code for instances of type_alias instances that are not unions."
                    )

            else:
                raise RuntimeError(
                    f"Cannot generate code for instances of kind '{type_['kind']}'"
                )

        elif kind == "dictionary_of":
            key_type, _ = self.get_python_type(p["type"]["key"])
            if "InstrumentedField" in key_type:
                value_type, _ = self.get_python_type(p["type"]["value"])
                if p["type"]["singleKey"]:
                    # special handling for single-key dicts with field key
                    k["args"] = [
                        {
                            "name": "_field",
                            "type": add_not_set(key_type),
                            "doc": [":arg _field: The field to use in this query."],
                            "required": False,
                            "positional": True,
                        },
                        {
                            "name": "_value",
                            "type": add_not_set(add_dict_type(value_type)),
                            "doc": [":arg _value: The query value for the field."],
                            "required": False,
                            "positional": True,
                        },
                    ]
                    k["is_single_field"] = True
                else:
                    # special handling for multi-key dicts with field keys
                    k["args"] = [
                        {
                            "name": "_fields",
                            "type": f"Optional[Mapping[{key_type}, {value_type}]]",
                            "doc": [
                                ":arg _fields: A dictionary of fields with their values."
                            ],
                            "required": False,
                            "positional": True,
                        },
                    ]
                    k["is_multi_field"] = True
            else:
                raise RuntimeError(f"Cannot generate code for type {p['type']}")

        else:
            raise RuntimeError(f"Cannot generate code for type {p['type']}")
        return [k] + other_classes

    def interface_to_python_class(
        self,
        interface,
        namespace=None,
        *,
        for_types_py=True,
        for_response=False,
    ):
        """Return a dictionary with template data necessary to render an
        interface a Python class.

        This is used to render interfaces that are referenced by container
        classes. The current list of rendered interfaces is passed as a second
        argument to allow this method to add more interfaces to it as they are
        discovered.

        The returned format is as follows:

        ```python
        {
            "name": "the class name to use for the interface class",
            "parent": "the parent class name",
            "args": [ # a Python description of each class attribute
                "name": "the name of the attribute",
                "type": "the Python type hint for the attribute",
                "doc": ["formatted lines of documentation to add to class docstring"],
                "required": bool,
                "positional": bool,
            ],
            "buckets_as_dict": "type" # optional, only present in aggregation response
                                      # classes that have buckets that can have a list
                                      # or dict representation
        }
        ```
        """
        type_ = self.find_type(interface, namespace)
        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": []}
        k["docstring"] = wrapped_doc(type_.get("description") or "")
        self.add_behaviors(
            type_, k, for_types_py=for_types_py, for_response=for_response
        )
        generics = []
        while True:
            for arg in type_["properties"]:
                if interface == "ResponseBody" and arg["name"] == "hits":
                    k["args"].append(
                        {
                            "name": "hits",
                            "type": "Sequence[_R]",
                            "doc": [":arg hits: search results"],
                            "required": arg["required"],
                        }
                    )
                elif interface == "ResponseBody" and arg["name"] == "aggregations":
                    # Aggregations are tricky because the DSL client uses a
                    # flexible representation that is difficult to generate
                    # from the schema.
                    # To handle this we let the generator do its work by calling
                    # `add_attribute()`, but then we save the generated attribute
                    # apart and replace it with the DSL's `AggResponse` class.
                    # The generated type is then used in type hints in variables
                    # and methods of this class.
                    self.add_attribute(
                        k, arg, for_types_py=for_types_py, for_response=for_response
                    )
                    k["aggregate_type"] = (
                        k["args"][-1]["type"]
                        .split("Mapping[str, ")[1]
                        .rsplit("]", 1)[0]
                    )
                    k["args"][-1] = {
                        "name": "aggregations",
                        "type": '"AggResponse[_R]"',
                        "doc": [":arg aggregations: aggregation results"],
                        "required": arg["required"],
                    }
                elif (
                    "name" in type_
                    and type_["name"]["name"] == "MultiBucketAggregateBase"
                    and arg["name"] == "buckets"
                ):
                    # Also during aggregation response generation, the "buckets"
                    # attribute that many aggregation responses have is very
                    # complex, supporting over a dozen different aggregation
                    # types via generics, each in array or object configurations.
                    # Typing this attribute proved very difficult. A solution
                    # that worked with mypy and pyright is to type "buckets"
                    # for the list form, and create a `buckets_as_dict`
                    # property that is typed appropriately for accessing the
                    # buckets in dictionary form.
                    # The generic type is assumed to be the first in the list,
                    # which is a simplification that should be improved when a
                    # more complete implementation of generics is added.
                    if generics[0]["type"]["name"] == "Void":
                        generic_type = "Any"
                    else:
                        _g = self.find_type(
                            generics[0]["type"]["name"],
                            generics[0]["type"]["namespace"],
                        )
                        generic_type, _ = self.get_python_type(
                            _g, for_response=for_response
                        )
                        generic_type = type_for_types_py(generic_type)
                    k["args"].append(
                        {
                            "name": arg["name"],
                            # for the type we only include the array form, since
                            # this client does not request the dict form
                            "type": f"Sequence[{generic_type}]",
                            "doc": [
                                ":arg buckets: (required) the aggregation buckets as a list"
                            ],
                            "required": True,
                        }
                    )
                    k["buckets_as_dict"] = generic_type
                else:
                    if interface == "Hit" and arg["name"].startswith("_"):
                        # Python DSL removes the undersore prefix from all the
                        # properties of the hit, so we do the same
                        arg["name"] = arg["name"][1:]

                    self.add_attribute(
                        k, arg, for_types_py=for_types_py, for_response=for_response
                    )

            if "inherits" not in type_ or "type" not in type_["inherits"]:
                break

            if "generics" in type_["inherits"]:
                # Generics are only supported for certain specific cases at this
                # time. Here we just save them so that they can be recalled later
                # while traversing over to parent classes to find inherited
                # attributes.
                for generic_type in type_["inherits"]["generics"]:
                    generics.append(generic_type)

            type_ = self.find_type(
                type_["inherits"]["type"]["name"],
                type_["inherits"]["type"]["namespace"],
            )
        return k


def generate_query_py(schema, filename):
    """Generate query.py with all the properties of `QueryContainer` as Python
    classes.
    """
    classes = []
    query_container = schema.find_type("QueryContainer", "_types.query_dsl")
    for p in query_container["properties"]:
        classes += schema.property_to_python_class(p)

    with open(filename, "wt") as f:
        f.write(query_py.render(classes=classes, parent="Query"))
    print(f"Generated {filename}.")


def generate_aggs_py(schema, filename):
    """Generate aggs.py with all the properties of `AggregationContainer` as
    Python classes.
    """
    classes = []
    aggs_container = schema.find_type("AggregationContainer", "_types.aggregations")
    for p in aggs_container["properties"]:
        if "containerProperty" not in p or not p["containerProperty"]:
            classes += schema.property_to_python_class(p)

    with open(filename, "wt") as f:
        f.write(aggs_py.render(classes=classes, parent="Agg"))
    print(f"Generated {filename}.")


def generate_response_init_py(schema, filename):
    """Generate response/__init__.py with all the response properties
    documented and typed.
    """
    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=search_response, ubq_response=ubq_response)
        )
    print(f"Generated {filename}.")


def generate_types_py(schema, filename):
    """Generate types.py"""
    classes = {}
    for interface in schema.interfaces:
        if interface == "PipeSeparatedFlags":
            continue  # handled as a special case
        for_response = interface in schema.response_interfaces
        k = schema.interface_to_python_class(
            interface, for_types_py=True, for_response=for_response
        )
        classes[k["name"]] = k

    # sort classes by being request/response and then by name
    sorted_classes = sorted(
        list(classes.keys()),
        key=lambda i: str(int(i in schema.response_interfaces)) + i,
    )
    classes_list = []
    for n in sorted_classes:
        k = classes[n]
        if k in classes_list:
            continue
        classes_list.append(k)

    with open(filename, "wt") as f:
        f.write(types_py.render(classes=classes_list))
    print(f"Generated {filename}.")


if __name__ == "__main__":
    schema = ElasticsearchSchema()
    generate_query_py(schema, "elasticsearch_dsl/query.py")
    generate_aggs_py(schema, "elasticsearch_dsl/aggs.py")
    generate_response_init_py(schema, "elasticsearch_dsl/response/__init__.py")
    # generate_response_hit_py(schema, "elasticsearch_dsl/response/hit.py")
    # generate_response_aggs_py(schema, "elasticsearch_dsl/response/aggs.py")
    generate_types_py(schema, "elasticsearch_dsl/types.py")