PY2 = sys.version_info[0] == 2
def force_unicode(value):
def with_metaclass(meta, *bases):
        def __new__(cls, name, this_bases, d):
    class metaclass(meta): PY2 = sys.version_info[0] == 2
def force_unicode(value):
def with_metaclass(meta, *bases):
        def __new__(cls, name, this_bases, d):
    class metaclass(meta): def _with_clone(fn):
    @wraps(fn)
    def wrapper(self, *args, **kwargs):
        clone = self.clone()
        res = fn(clone, *args, **kwargs)
        if res is not None:
            return res
        return clone
    return wrapper def merge_params(params, args, kwargs):
    assert isinstance(args, Iterable), args
    assert isinstance(kwargs, Mapping), kwargs
    new = dict()
    for a in args:
        new.update(a)
    new.update(kwargs)
    return type(params)(params, **new) def collect_doc_classes(expr):
    if hasattr(expr, '_collect_doc_classes'):
        return expr._collect_doc_classes()

    if isinstance(expr, dict):
        return set().union(
            *[collect_doc_classes(e)
              for e in chain(expr.keys(), expr.values())]
        )

    if isinstance(expr, (list, tuple)):
        return set().union(*[collect_doc_classes(e) for e in expr])

    return set() class Params(Expression, Mapping):
    __visit_name__ = 'params'

    def __init__(self, *args, **kwargs):
        params = {}
        for d in args:
            params.update(d)
        params.update(kwargs)
        self._params = {}
        for k, v in clean_params(params).items():
            if k.endswith('_') and not k.startswith('_'):
                k = k.rstrip('_')
            self._params[k] = v

    def _collect_doc_classes(self):
        return collect_doc_classes(self._params)

    def __len__(self):
        return len(self._params)

    def __iter__(self):
        return iter(self._params)

    def __getitem__(self, key):
        return self._params[key]

    def __contains__(self, key):
        return key in self._params class Source(Expression):
    __visit_name__ = 'source'

    def __init__(self, fields, include=None, exclude=None):
        self.fields = fields
        self.include = include
        self.exclude = exclude

    def _collect_doc_classes(self):
        return set().union(
            collect_doc_classes(self.fields),
            collect_doc_classes(self.include),
            collect_doc_classes(self.exclude),
        ) class Highlight(Expression):
    __visit_name__ = 'highlight'

    def __init__(self, fields=None, **kwargs):
        self.fields = fields
        self.params = Params(kwargs)

    def _collect_doc_classes(self):
        return set().union(
            collect_doc_classes(self.fields),
            collect_doc_classes(self.params),
        ) class Rescore(Expression):
    __visit_name__ = 'rescore'

    def __init__(self, rescorer, window_size=None,
                 ):
        self.rescorer = rescorer
        self.window_size = window_size

    def _collect_doc_classes(self):
        return collect_doc_classes(self.rescorer) class Script(Expression):
    __visit_name__ = 'script'

    def __init__(self, inline=None, lang=None,
                 params=None, id=None, file=None):
        self.inline = inline
        self.id = id
        self.lang = lang
        self.file = file
        self.script_params = params
import warnings
from abc import ABCMeta
from collections import namedtuple, OrderedDict
from .compat import zip, with_metaclass, string_types
from .compat import Iterable
from .util import _with_clone
from .util import merge_params, collect_doc_classes
from .expression import Params, Source, Highlight, Rescore, Script
           elasticsearch. Can be one of the following types:

           - field expression, for example: ``PostDocument.title``
           - ``str`` means field name or glob pattern. For example:
             ``"title"``, ``"user.*"``
           - ``False`` disables retrieving source
           - ``True`` enables retrieving all source document
           - ``None`` cancels source filtering applied before

        :param include: list of fields to include

        :param exclude: list of fields to exclude

        Example:

        .. testcode:: source

           search_query = SearchQuery().source(PostDocument.title, 'user.*')

        .. testcode:: source

           assert search_query.to_dict(Compiler_5_0) == {'_source': ['title', 'user.*']}

        See `source filtering <https://www.elastic.co/guide/en/elasticsearch/reference/current/search-request-source-filtering.html>`_
        for more information.
        """  # noqa:E501
        if len(fields) == 1 and fields[0] is None:
            if '_source' in self.__dict__:
                del self._source
        elif len(fields) == 1 and isinstance(fields[0], bool):
            self._source = Source(fields[0], **kwargs)
        else:
            self._source = Source(fields, **kwargs)

    @_with_clone
    def stored_fields(self, *fields):
        """Allows to load fields that marked as ``store: true``.

        Example:

        .. testcode:: stored_fields

           search_query = SearchQuery().stored_fields(PostDocument.rank)

        .. testcode:: stored_fields

           assert search_query.to_dict(Compiler_5_0) == {'stored_fields': ['rank']}

        See `stored fields <https://www.elastic.co/guide/en/elasticsearch/reference/current/mapping-store.html>`_ and
        `stored fields filtering <https://www.elastic.co/guide/en/elasticsearch/reference/current/search-request-stored-fields.html>`_
        for more information.
        """  # noqa:E501
        if len(fields) == 1 and fields[0] is None:
            if '_fields' in self.__dict__:
                del self._fields
        elif len(fields) == 1 and isinstance(fields[0], bool):
            self._fields = fields[0]
        else:
            self._fields = fields

    def fields(self, *fields):
        return self.stored_fields(*fields)

    @_with_clone
    def script_fields(self, *args, **kwargs):
        """Allows to evaluate fields based on scripts.

        .. testcode:: script_fields

           from elasticmagic import Script

           search_query = SearchQuery().script_fields(
               rating=Script(
                   inline='doc[params.positive_opinions_field].value / '
                          'doc[params.total_opinions_field].value * 5',
                   params={
                       'total_opinions_field': PostDocument.total_opinions,
                       'positive_opinions_field': PostDocument.positive_opinions,
                   }
               )
           )

        .. testcode:: script_fields

           expected = {
               'script_fields': {
                   'rating': {
                       'script': {
                           'inline': 'doc[params.positive_opinions_field].value / '
                                     'doc[params.total_opinions_field].value * 5',
                           'params': {
                               'total_opinions_field': 'total_opinions',
                               'positive_opinions_field': 'positive_opinions'
                           }
                       }
                   }
               }
           }

           assert search_query.to_dict(Compiler_5_0) == {
               'script_fields': {
                   'rating': {
                       'script': {
                           'inline': 'doc[params.positive_opinions_field].value / '
                                     'doc[params.total_opinions_field].value * 5',
                           'params': {
                               'total_opinions_field': 'total_opinions',
                               'positive_opinions_field': 'positive_opinions'
                           }
                       }
                   }
               }
           }

        See `script fields <https://www.elastic.co/guide/en/elasticsearch/reference/current/search-request-script-fields.html>`_
        """  # noqa:E501
        if len(args) == 1 and args[0] is None:
            if '_script_fields' in self.__dict__:
                del self._script_fields
        for key in kwargs:

            if isinstance(kwargs[key], Script):