<commit_before>from __future__ import absolute_import

from django.db import models

from elasticsearch import Elasticsearch, NotFoundError, RequestError
from elasticsearch.helpers import bulk

from wagtail.wagtailsearch.backends.base import BaseSearch
from wagtail.wagtailsearch.indexed import Indexed

import string


class ElasticSearchQuery(object):
    def __init__(self, model, query_string, fields=None, filters={}):
        self.model = model
        self.query_string = query_string
        self.fields = fields or ['_all']
        self.filters = filters

    def _get_filters(self):
        # Filters
        filters = []

        # Filter by content type
        filters.append({
            'prefix': {
                'content_type': self.model.indexed_get_content_type()
            }
        })

        # Extra filters
        if self.filters:
            for key, value in self.filters.items():
                if '__' in key:
                    field, lookup = key.split('__')
                else:
                    field = key
                    lookup = None

                if lookup is None:
                    if value is None:
                        filters.append({
                            'missing': {
                                'field': field,
                            }
                        })
                    else:
                        filters.append({
                            'term': {
                                field: value
                            }
                        })

                if lookup in ['startswith', 'prefix']:
                    filters.append({
                        'prefix': {
                            field: value
                        }
                    })

                if lookup in ['gt', 'gte', 'lt', 'lte']:
                    filters.append({
                        'range': {
                            field: {
                                lookup: value,
                            }
                        }
                    })

                if lookup == 'range':
                    lower, upper = value
                    filters.append({
                        'range': {
                            field: {
                                'gte': lower,
                                'lte': upper,
                            }
                        }
                    })

        return filters

    def to_es(self):
        # Query
        query = {
            'query_string': {
                'query': self.query_string,
            }
        }

        # Fields
        if self.fields:
            query['query_string']['fields'] = self.fields

        # Filters
        filters = self._get_filters()

        return {
            'filtered': {
                'query': query,
                'filter': {
                    'and': filters,
                }
            }
        }


class ElasticSearchResults(object):
    def __init__(self, backend, query, prefetch_related=[]):
        self.backend = backend
        self.query = query
        self.prefetch_related = prefetch_related

    def _get_results_pks(self, offset=0, limit=None):
        # Params for elasticsearch query
        params = dict(
            index=self.backend.es_index,
            body=dict(query=self.query.to_es()),
            _source=False,
            fields='pk',
            from_=offset,
        )

        # Add limit if set
        if limit is not None:
            params['size'] = limit

        # Send to ElasticSearch
        hits = self.backend.es.search(**params)

        pks = [hit['fields']['pk'] for hit in hits['hits']['hits']]

        # ElasticSearch 1.x likes to pack pks into lists, unpack them if this has happened
        return [pk[0] if isinstance(pk, list) else pk for pk in pks]

    def _get_count(self):
        query = self.query.to_es()

        # Elasticsearch 1.x
        count = self.backend.es.count(
            index=self.backend.es_index,
            body=dict(query=query),
        )

        # ElasticSearch 0.90.x fallback
        if not count['_shards']['successful'] and "No query registered for [query]]" in count['_shards']['failures'][0]['reason']:
            count = self.backend.es.count(
                index=self.backend.es_index,
                body=query,
            )

        return count['count']

    def __getitem__(self, key):
        if isinstance(key, slice):
            # Get primary keys
            pk_list_unclean = self._get_results_pks(key.start, key.stop - key.start)

            # Remove duplicate keys (and preserve order)
            seen_pks = set()
            pk_list = []
            for pk in pk_list_unclean:
                if pk not in seen_pks:
                    seen_pks.add(pk)
                    pk_list.append(pk)

            # Get results
            results = self.query.model.objects.filter(pk__in=pk_list)

            # Prefetch related
            for prefetch in self.prefetch_related:
                results = results.prefetch_related(prefetch)

            # Put results into a dictionary (using primary key as the key)
            results_dict = dict((str(result.pk), result) for result in results)

            # Build new list with items in the correct order
            results_sorted = [results_dict[str(pk)] for pk in pk_list if str(pk) in results_dict]

            # Return the list
            return results_sorted
        else:
            # Return a single item
            pk = self._get_results_pks(key, key + 1)[0]
            return self.query.model.objects.get(pk=pk)

    def __len__(self):
        return self._get_count()


class ElasticSearch(BaseSearch):
    def __init__(self, params):
        super(ElasticSearch, self).__init__(params)

        # Get settings
        self.es_urls = params.pop('URLS', ['http://localhost:9200'])
        self.es_index = params.pop('INDEX', 'wagtail')
        self.es_timeout = params.pop('TIMEOUT', 5)
        self.es_force_new = params.pop('FORCE_NEW', False)

        # Get ElasticSearch interface
        # Any remaining params are passed into the ElasticSearch constructor
        self.es = Elasticsearch(
            urls=self.es_urls,
            timeout=self.es_timeout,
            force_new=self.es_force_new,
            **params)

    def reset_index(self):
        # Delete old index
        try:
            self.es.indices.delete(self.es_index)
        except NotFoundError:
            pass

        # Settings
        INDEX_SETTINGS = {<commit_msg>Renamed a couple of methods on ElasticSearchResults<commit_after>from __future__ import absolute_import

from django.db import models

from elasticsearch import Elasticsearch, NotFoundError, RequestError
from elasticsearch.helpers import bulk

from wagtail.wagtailsearch.backends.base import BaseSearch
from wagtail.wagtailsearch.indexed import Indexed

import string


class ElasticSearchQuery(object):
    def __init__(self, model, query_string, fields=None, filters={}):
        self.model = model
        self.query_string = query_string
        self.fields = fields or ['_all']
        self.filters = filters

    def _get_filters(self):
        # Filters
        filters = []

        # Filter by content type
        filters.append({
            'prefix': {
                'content_type': self.model.indexed_get_content_type()
            }
        })

        # Extra filters
        if self.filters:
            for key, value in self.filters.items():
                if '__' in key:
                    field, lookup = key.split('__')
                else:
                    field = key
                    lookup = None

                if lookup is None:
                    if value is None:
                        filters.append({
                            'missing': {
                                'field': field,
                            }
                        })
                    else:
                        filters.append({
                            'term': {
                                field: value
                            }
                        })

                if lookup in ['startswith', 'prefix']:
                    filters.append({
                        'prefix': {
                            field: value
                        }
                    })

                if lookup in ['gt', 'gte', 'lt', 'lte']:
                    filters.append({
                        'range': {
                            field: {
                                lookup: value,
                            }
                        }
                    })

                if lookup == 'range':
                    lower, upper = value
                    filters.append({
                        'range': {
                            field: {
                                'gte': lower,
                                'lte': upper,
                            }
                        }
                    })

        return filters

    def to_es(self):
        # Query
        query = {
            'query_string': {
                'query': self.query_string,
            }
        }

        # Fields
        if self.fields:
            query['query_string']['fields'] = self.fields

        # Filters
        filters = self._get_filters()

        return {
            'filtered': {
                'query': query,
                'filter': {
                    'and': filters,
                }
            }
        }


class ElasticSearchResults(object):
    def __init__(self, backend, query, prefetch_related=[]):
        self.backend = backend
        self.query = query
        self.prefetch_related = prefetch_related

    def _do_search(self, offset=0, limit=None):
        # Params for elasticsearch query
        params = dict(
            index=self.backend.es_index,
            body=dict(query=self.query.to_es()),
            _source=False,
            fields='pk',
            from_=offset,
        )

        # Add limit if set
        if limit is not None:
            params['size'] = limit

        # Send to ElasticSearch
        hits = self.backend.es.search(**params)

        pks = [hit['fields']['pk'] for hit in hits['hits']['hits']]

        # ElasticSearch 1.x likes to pack pks into lists, unpack them if this has happened
        return [pk[0] if isinstance(pk, list) else pk for pk in pks]

    def _do_count(self):
        query = self.query.to_es()

        # Elasticsearch 1.x
        count = self.backend.es.count(
            index=self.backend.es_index,
            body=dict(query=query),
        )

        # ElasticSearch 0.90.x fallback
        if not count['_shards']['successful'] and "No query registered for [query]]" in count['_shards']['failures'][0]['reason']:
            count = self.backend.es.count(
                index=self.backend.es_index,
                body=query,
            )

        return count['count']

    def __getitem__(self, key):
        if isinstance(key, slice):
            # Get primary keys
            pk_list_unclean = self._do_search(key.start, key.stop - key.start)

            # Remove duplicate keys (and preserve order)
            seen_pks = set()
            pk_list = []
            for pk in pk_list_unclean:
                if pk not in seen_pks:
                    seen_pks.add(pk)
                    pk_list.append(pk)

            # Get results
            results = self.query.model.objects.filter(pk__in=pk_list)

            # Prefetch related
            for prefetch in self.prefetch_related:
                results = results.prefetch_related(prefetch)

            # Put results into a dictionary (using primary key as the key)
            results_dict = dict((str(result.pk), result) for result in results)

            # Build new list with items in the correct order
            results_sorted = [results_dict[str(pk)] for pk in pk_list if str(pk) in results_dict]

            # Return the list
            return results_sorted
        else:
            # Return a single item
            pk = self._do_search(key, key + 1)[0]
            return self.query.model.objects.get(pk=pk)

    def __len__(self):
        return self._do_count()


class ElasticSearch(BaseSearch):
    def __init__(self, params):
        super(ElasticSearch, self).__init__(params)

        # Get settings
        self.es_urls = params.pop('URLS', ['http://localhost:9200'])
        self.es_index = params.pop('INDEX', 'wagtail')
        self.es_timeout = params.pop('TIMEOUT', 5)
        self.es_force_new = params.pop('FORCE_NEW', False)

        # Get ElasticSearch interface
        # Any remaining params are passed into the ElasticSearch constructor
        self.es = Elasticsearch(
            urls=self.es_urls,
            timeout=self.es_timeout,
            force_new=self.es_force_new,
            **params)

    def reset_index(self):
        # Delete old index
        try:
            self.es.indices.delete(self.es_index)
        except NotFoundError:
            pass

        # Settings
        INDEX_SETTINGS = {
{'fraction_non_alphanumeric': 0.09063260340632603, 'fraction_numerical': 0.002281021897810219, 'mean_word_length': 1.2966649205517724, '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': '7479735', 'n_tokens_mistral': 3493, 'n_tokens_neox': 3321, 'n_words': 828}