<commit_before>from .utils import EqualityComparableUsingAttributeDictionary
from .filters import Filter, TermFilter, TermsFilter, ANDFilter, NotFilter


class FacetFactory(EqualityComparableUsingAttributeDictionary):

    def __init__(self):
        self.facets = []

    def add_term_facet(self, *args, **kwargs):
        """Add a term factory facet"""
        self.facets.append(TermFacet(*args, **kwargs))

    def add_date_facet(self, *args, **kwargs):
        """Add a date factory facet"""
        self.facets.append(DateHistogramFacet(*args, **kwargs))

    def add_geo_facet(self, *args, **kwargs):
        """Add a geo factory facet"""
        self.facets.append(GeoDistanceFacet(*args, **kwargs))

    def add(self, facet):
        """Add a term factory"""
        self.facets.append(facet)

    def reset(self):
        """Reset the facets"""
        self.facets = []

    def serialize(self):
        res = {}
        for facet in self.facets:
            res.update(facet.serialize())
        return res


class Facet(EqualityComparableUsingAttributeDictionary):

    def __init__(self, name, scope=None, nested=None, is_global=None,
                 facet_filter=None, **kwargs):
        self.name = name
        self.scope = scope
        self.nested = nested
        self.is_global = is_global
        self.facet_filter = facet_filter

    def serialize(self):
        data = {self._internal_name: self._serialize()}
        if self.scope is not None:
            data["scope"] = self.scope
        if self.nested is not None:
            data["nested"] = self.nested
        if self.is_global:
            data['global'] = self.is_global
        if self.facet_filter:
            data['facet_filter'] = self.facet_filter.serialize()
        return {self.name: data}

    def _serialize(self):
        raise NotImplementedError

    @property
    def _internal_name(self):
        raise NotImplementedError


# TODO: remove these
FacetFilter = Filter
TermFacetFilter = TermFilter
TermsFacetFilter = TermsFilter
ANDFacetFilter = ANDFilter
NotFacetFilter = NotFilter


class QueryFacet(Facet):

    _internal_name = "query"

    def __init__(self, name, query, **kwargs):
        super(QueryFacet, self).__init__(name, **kwargs)
        self.query = query

    def _serialize(self):
        return self.query.serialize()


class FilterFacet(Facet):

    _internal_name = "filter"

    def __init__(self, name, filter, **kwargs):
        super(FilterFacet, self).__init__(name, **kwargs)
        self.filter = filter

    def _serialize(self):
        return self.filter.serialize()


class HistogramFacet(Facet):

    _internal_name = "histogram"

    def __init__(self, name, field=None, interval=None, time_interval=None,
                 key_field=None, value_field=None, key_script=None,
                 value_script=None, params=None, **kwargs):
        super(HistogramFacet, self).__init__(name, **kwargs)
        self.field = field
        self.interval = interval
        self.time_interval = time_interval
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _add_interval(self, data):
        if self.interval:
            data['interval'] = self.interval
        elif self.time_interval:
            data['time_interval'] = self.time_interval
        else:
            raise RuntimeError("Invalid field: interval or time_interval required")

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
            self._add_interval(data)
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            else:
                raise RuntimeError("Invalid key_field: value_field required")
            self._add_interval(data)
        elif self.key_script:
            data['key_script'] = self.key_script
            if self.value_script:
                data['value_script'] = self.value_script
            else:
                raise RuntimeError("Invalid key_script: value_script required")
            if self.params:
                data['params'] = self.params
            if self.interval:
                data['interval'] = self.interval
            elif self.time_interval:
                data['time_interval'] = self.time_interval
        return data


class DateHistogramFacet(Facet):

    _internal_name = "date_histogram"

    def __init__(self, name, field=None, interval=None, zone=None, key_field=None,
                 value_field=None, value_script=None, params=None, **kwargs):
        super(DateHistogramFacet, self).__init__(name, **kwargs)
        self.field = field
        self.interval = interval
        self.zone = zone
        self.key_field = key_field
        self.value_field = value_field
        self.value_script = value_script
        self.params = params

    def _serialize(self):
        data = {}
        if self.interval:
            data['interval'] = self.interval
            if self.zone:
                data['zone'] = self.zone
        else:
            raise RuntimeError("interval required")
        if self.field:
            data['field'] = self.field
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            elif self.value_script:
                data['value_script'] = self.value_script
                if self.params:
                    data['params'] = self.params
            else:
                raise RuntimeError("Invalid key_field: value_field or value_script required")
        return data


class RangeFacet(Facet):

    _internal_name = "range"

    def __init__(self, name, field=None, ranges=None, key_field=None, value_field=None,
                 key_script=None, value_script=None, params=None, **kwargs):
        super(RangeFacet, self).__init__(name, **kwargs)
        self.field = field
        self.ranges = ranges or []
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _serialize(self):<commit_msg>fix DateHistogramFacet reference from 'zone' to 'time_zone' and add additional params<commit_after>from .utils import EqualityComparableUsingAttributeDictionary
from .filters import Filter, TermFilter, TermsFilter, ANDFilter, NotFilter


class FacetFactory(EqualityComparableUsingAttributeDictionary):

    def __init__(self):
        self.facets = []

    def add_term_facet(self, *args, **kwargs):
        """Add a term factory facet"""
        self.facets.append(TermFacet(*args, **kwargs))

    def add_date_facet(self, *args, **kwargs):
        """Add a date factory facet"""
        self.facets.append(DateHistogramFacet(*args, **kwargs))

    def add_geo_facet(self, *args, **kwargs):
        """Add a geo factory facet"""
        self.facets.append(GeoDistanceFacet(*args, **kwargs))

    def add(self, facet):
        """Add a term factory"""
        self.facets.append(facet)

    def reset(self):
        """Reset the facets"""
        self.facets = []

    def serialize(self):
        res = {}
        for facet in self.facets:
            res.update(facet.serialize())
        return res


class Facet(EqualityComparableUsingAttributeDictionary):

    def __init__(self, name, scope=None, nested=None, is_global=None,
                 facet_filter=None, **kwargs):
        self.name = name
        self.scope = scope
        self.nested = nested
        self.is_global = is_global
        self.facet_filter = facet_filter

    def serialize(self):
        data = {self._internal_name: self._serialize()}
        if self.scope is not None:
            data["scope"] = self.scope
        if self.nested is not None:
            data["nested"] = self.nested
        if self.is_global:
            data['global'] = self.is_global
        if self.facet_filter:
            data['facet_filter'] = self.facet_filter.serialize()
        return {self.name: data}

    def _serialize(self):
        raise NotImplementedError

    @property
    def _internal_name(self):
        raise NotImplementedError


# TODO: remove these
FacetFilter = Filter
TermFacetFilter = TermFilter
TermsFacetFilter = TermsFilter
ANDFacetFilter = ANDFilter
NotFacetFilter = NotFilter


class QueryFacet(Facet):

    _internal_name = "query"

    def __init__(self, name, query, **kwargs):
        super(QueryFacet, self).__init__(name, **kwargs)
        self.query = query

    def _serialize(self):
        return self.query.serialize()


class FilterFacet(Facet):

    _internal_name = "filter"

    def __init__(self, name, filter, **kwargs):
        super(FilterFacet, self).__init__(name, **kwargs)
        self.filter = filter

    def _serialize(self):
        return self.filter.serialize()


class HistogramFacet(Facet):

    _internal_name = "histogram"

    def __init__(self, name, field=None, interval=None, time_interval=None,
                 key_field=None, value_field=None, key_script=None,
                 value_script=None, params=None, **kwargs):
        super(HistogramFacet, self).__init__(name, **kwargs)
        self.field = field
        self.interval = interval
        self.time_interval = time_interval
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _add_interval(self, data):
        if self.interval:
            data['interval'] = self.interval
        elif self.time_interval:
            data['time_interval'] = self.time_interval
        else:
            raise RuntimeError("Invalid field: interval or time_interval required")

    def _serialize(self):
        data = {}
        if self.field:
            data['field'] = self.field
            self._add_interval(data)
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            else:
                raise RuntimeError("Invalid key_field: value_field required")
            self._add_interval(data)
        elif self.key_script:
            data['key_script'] = self.key_script
            if self.value_script:
                data['value_script'] = self.value_script
            else:
                raise RuntimeError("Invalid key_script: value_script required")
            if self.params:
                data['params'] = self.params
            if self.interval:
                data['interval'] = self.interval
            elif self.time_interval:
                data['time_interval'] = self.time_interval
        return data


class DateHistogramFacet(Facet):

    _internal_name = "date_histogram"

    def __init__(self, name, field=None, interval=None, time_zone=None, pre_zone=None,
                 post_zone=None, factor=None, pre_offset=None, post_offset=None,
                 key_field=None, value_field=None, value_script=None, params=None, **kwargs):
        super(DateHistogramFacet, self).__init__(name, **kwargs)
        self.field = field
        self.interval = interval
        self.time_zone = time_zone
        self.pre_zone = pre_zone
        self.post_zone = post_zone
        self.factor = factor
        self.pre_offset = pre_offset
        self.post_offset = post_offset
        self.key_field = key_field
        self.value_field = value_field
        self.value_script = value_script
        self.params = params

    def _serialize(self):
        data = {}
        if self.interval:
            data['interval'] = self.interval
        else:
            raise RuntimeError("interval required")
        if self.time_zone:
            data['time_zone'] = self.time_zone
        if self.pre_zone:
            data['pre_zone'] = self.pre_zone
        if self.post_zone:
            data['post_zone'] = self.post_zone
        if self.factor:
            data['factor'] = self.factor
        if self.pre_offset:
            data['pre_offset'] = self.pre_offset
        if self.post_offset:
            data['post_offset'] = self.post_offset
        if self.field:
            data['field'] = self.field
        elif self.key_field:
            data['key_field'] = self.key_field
            if self.value_field:
                data['value_field'] = self.value_field
            elif self.value_script:
                data['value_script'] = self.value_script
                if self.params:
                    data['params'] = self.params
            else:
                raise RuntimeError("Invalid key_field: value_field or value_script required")
        return data


class RangeFacet(Facet):

    _internal_name = "range"

    def __init__(self, name, field=None, ranges=None, key_field=None, value_field=None,
                 key_script=None, value_script=None, params=None, **kwargs):
        super(RangeFacet, self).__init__(name, **kwargs)
        self.field = field
        self.ranges = ranges or []
        self.key_field = key_field
        self.value_field = value_field
        self.key_script = key_script
        self.value_script = value_script
        self.params = params

    def _serialize(self):
{'fraction_non_alphanumeric': 0.09912189313885995, 'fraction_numerical': 0.0, 'mean_word_length': 2.2842130987292277, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2828434', 'n_tokens_mistral': 3836, 'n_tokens_neox': 3726, 'n_words': 842}