<commit_before>from connection import _get_db

import pymongo
import copy


__all__ = ['queryset_manager', 'Q', 'InvalidQueryError', 
           'InvalidCollectionError']

# The maximum number of items to display in a QuerySet.__repr__
REPR_OUTPUT_SIZE = 20


class InvalidQueryError(Exception):
    pass


class OperationError(Exception):
    pass


class Q(object):
    
    OR = '||'
    AND = '&&'
    OPERATORS = {
        'eq': 'this.%(field)s == %(value)s',
        'neq': 'this.%(field)s != %(value)s',
        'gt': 'this.%(field)s > %(value)s',
        'gte': 'this.%(field)s >= %(value)s',
        'lt': 'this.%(field)s < %(value)s',
        'lte': 'this.%(field)s <= %(value)s',
        'lte': 'this.%(field)s <= %(value)s',
        'in': 'this.%(field)s.indexOf(%(value)s) != -1',
        'nin': 'this.%(field)s.indexOf(%(value)s) == -1',
        'mod': '%(field)s %% %(value)s',
        'all': ('%(value)s.every(function(a){'
                'return this.%(field)s.indexOf(a) != -1 })'),
        'size': 'this.%(field)s.length == %(value)s',
        'exists': 'this.%(field)s != null',
    }
    
    def __init__(self, **query):
        self.query = [query]

    def _combine(self, other, op):
        obj = Q()
        obj.query = ['('] + copy.deepcopy(self.query) + [op]
        obj.query += copy.deepcopy(other.query) + [')']
        return obj

    def __or__(self, other):
        return self._combine(other, self.OR)

    def __and__(self, other):
        return self._combine(other, self.AND)

    def as_js(self, document):
        js = []
        js_scope = {}
        for i, item in enumerate(self.query):
            if isinstance(item, dict):
                item_query = QuerySet._transform_query(document, **item)
                # item_query will values will either be a value or a dict
                js.append(self._item_query_as_js(item_query, js_scope, i))
            else:
                js.append(item)
        return pymongo.code.Code(' '.join(js), js_scope)

    def _item_query_as_js(self, item_query, js_scope, item_num):
        # item_query will be in one of the following forms
        #    {'age': 25, 'name': 'Test'}
        #    {'age': {'$lt': 25}, 'name': {'$in': ['Test', 'Example']}
        #    {'age': {'$lt': 25, '$gt': 18}}
        js = []
        for i, (key, value) in enumerate(item_query.items()):
            op = 'eq'
            # Construct a variable name for the value in the JS
            value_name = 'i%sf%s' % (item_num, i)
            if isinstance(value, dict):
                # Multiple operators for this field
                for j, (op, value) in enumerate(value.items()):
                    # Create a custom variable name for this operator
                    op_value_name = '%so%s' % (value_name, j)
                    # Update the js scope with the value for this op
                    js_scope[op_value_name] = value
                    # Construct the JS that uses this op
                    operation_js = Q.OPERATORS[op.strip('$')] % {
                        'field': key, 
                        'value': op_value_name
                    }
                    js.append(operation_js)
            else:
                js_scope[value_name] = value
                # Construct the JS for this field
                field_js = Q.OPERATORS[op.strip('$')] % {
                    'field': key, 
                    'value': value_name
                }
                js.append(field_js)
        return ' && '.join(js)


class QuerySet(object):
    """A set of results returned from a query. Wraps a MongoDB cursor, 
    providing :class:`~mongoengine.Document` objects as the results.
    """
    
    def __init__(self, document, collection):
        self._document = document
        self._collection_obj = collection
        self._accessed_collection = False
        self._query = {}
        self._where_clauses = []

        # If inheritance is allowed, only return instances and instances of
        # subclasses of the class being used
        if document._meta.get('allow_inheritance'):
            self._query = {'_types': self._document._class_name}
        self._cursor_obj = None
        
    def ensure_index(self, key_or_list):
        """Ensure that the given indexes are in place.

        :param key_or_list: a single index key or a list of index keys (to
            construct a multi-field index); keys may be prefixed with a **+**
            or a **-** to determine the index ordering
        """
        if isinstance(key_or_list, basestring):
            key_or_list = [key_or_list]

        index_list = []
        # If _types is being used, prepend it to every specified index
        if self._document._meta.get('allow_inheritance'):
            index_list.append(('_types', 1))

        for key in key_or_list:
            # Get direction from + or -
            direction = pymongo.ASCENDING
            if key.startswith("-"):
                direction = pymongo.DESCENDING
            if key.startswith(("+", "-")):
                    key = key[1:]
            # Use real field name
            key = QuerySet._translate_field_name(self._document, key)
            index_list.append((key, direction))
        self._collection.ensure_index(index_list)
        return self

    def __call__(self, *q_objs, **query):
        """Filter the selected documents by calling the 
        :class:`~mongoengine.QuerySet` with a query.

        :param q_objs: :class:`~mongoengine.Q` objects to be used in the query
        :param query: Django-style query keyword arguments
        """
        for q in q_objs:
            self._where_clauses.append(q.as_js(self._document))
        query = QuerySet._transform_query(_doc_cls=self._document, **query)
        self._query.update(query)
        return self
        
    def filter(self, **query):
        return self.__call__(**query)

    @property
    def _collection(self):
        """Property that returns the collection object. This allows us to
        perform operations only if the collection is accessed.
        """
        if not self._accessed_collection:
            self._accessed_collection = True

            # Ensure document-defined indexes are created
            if self._document._meta['indexes']:
                for key_or_list in self._document._meta['indexes']:
                    self.ensure_index(key_or_list)

            # Ensure indexes created by uniqueness constraints
            for index in self._document._meta['unique_indexes']:
                self._collection.ensure_index(index, unique=True)

            # If _types is being used (for polymorphism), it needs an index
            if '_types' in self._query:
                self._collection.ensure_index('_types')
        return self._collection_obj

    @property
    def _cursor(self):
        if not self._cursor_obj:
            self._cursor_obj = self._collection.find(self._query)
            # Apply where clauses to cursor
            for js in self._where_clauses:
                self._cursor_obj.where(js)<commit_msg>added Q object support to QuerySet.filter<commit_after>from connection import _get_db

import pymongo
import copy


__all__ = ['queryset_manager', 'Q', 'InvalidQueryError', 
           'InvalidCollectionError']

# The maximum number of items to display in a QuerySet.__repr__
REPR_OUTPUT_SIZE = 20


class InvalidQueryError(Exception):
    pass


class OperationError(Exception):
    pass


class Q(object):
    
    OR = '||'
    AND = '&&'
    OPERATORS = {
        'eq': 'this.%(field)s == %(value)s',
        'neq': 'this.%(field)s != %(value)s',
        'gt': 'this.%(field)s > %(value)s',
        'gte': 'this.%(field)s >= %(value)s',
        'lt': 'this.%(field)s < %(value)s',
        'lte': 'this.%(field)s <= %(value)s',
        'lte': 'this.%(field)s <= %(value)s',
        'in': 'this.%(field)s.indexOf(%(value)s) != -1',
        'nin': 'this.%(field)s.indexOf(%(value)s) == -1',
        'mod': '%(field)s %% %(value)s',
        'all': ('%(value)s.every(function(a){'
                'return this.%(field)s.indexOf(a) != -1 })'),
        'size': 'this.%(field)s.length == %(value)s',
        'exists': 'this.%(field)s != null',
    }
    
    def __init__(self, **query):
        self.query = [query]

    def _combine(self, other, op):
        obj = Q()
        obj.query = ['('] + copy.deepcopy(self.query) + [op]
        obj.query += copy.deepcopy(other.query) + [')']
        return obj

    def __or__(self, other):
        return self._combine(other, self.OR)

    def __and__(self, other):
        return self._combine(other, self.AND)

    def as_js(self, document):
        js = []
        js_scope = {}
        for i, item in enumerate(self.query):
            if isinstance(item, dict):
                item_query = QuerySet._transform_query(document, **item)
                # item_query will values will either be a value or a dict
                js.append(self._item_query_as_js(item_query, js_scope, i))
            else:
                js.append(item)
        return pymongo.code.Code(' '.join(js), js_scope)

    def _item_query_as_js(self, item_query, js_scope, item_num):
        # item_query will be in one of the following forms
        #    {'age': 25, 'name': 'Test'}
        #    {'age': {'$lt': 25}, 'name': {'$in': ['Test', 'Example']}
        #    {'age': {'$lt': 25, '$gt': 18}}
        js = []
        for i, (key, value) in enumerate(item_query.items()):
            op = 'eq'
            # Construct a variable name for the value in the JS
            value_name = 'i%sf%s' % (item_num, i)
            if isinstance(value, dict):
                # Multiple operators for this field
                for j, (op, value) in enumerate(value.items()):
                    # Create a custom variable name for this operator
                    op_value_name = '%so%s' % (value_name, j)
                    # Update the js scope with the value for this op
                    js_scope[op_value_name] = value
                    # Construct the JS that uses this op
                    operation_js = Q.OPERATORS[op.strip('$')] % {
                        'field': key, 
                        'value': op_value_name
                    }
                    js.append(operation_js)
            else:
                js_scope[value_name] = value
                # Construct the JS for this field
                field_js = Q.OPERATORS[op.strip('$')] % {
                    'field': key, 
                    'value': value_name
                }
                js.append(field_js)
        return ' && '.join(js)


class QuerySet(object):
    """A set of results returned from a query. Wraps a MongoDB cursor, 
    providing :class:`~mongoengine.Document` objects as the results.
    """
    
    def __init__(self, document, collection):
        self._document = document
        self._collection_obj = collection
        self._accessed_collection = False
        self._query = {}
        self._where_clauses = []

        # If inheritance is allowed, only return instances and instances of
        # subclasses of the class being used
        if document._meta.get('allow_inheritance'):
            self._query = {'_types': self._document._class_name}
        self._cursor_obj = None
        
    def ensure_index(self, key_or_list):
        """Ensure that the given indexes are in place.

        :param key_or_list: a single index key or a list of index keys (to
            construct a multi-field index); keys may be prefixed with a **+**
            or a **-** to determine the index ordering
        """
        if isinstance(key_or_list, basestring):
            key_or_list = [key_or_list]

        index_list = []
        # If _types is being used, prepend it to every specified index
        if self._document._meta.get('allow_inheritance'):
            index_list.append(('_types', 1))

        for key in key_or_list:
            # Get direction from + or -
            direction = pymongo.ASCENDING
            if key.startswith("-"):
                direction = pymongo.DESCENDING
            if key.startswith(("+", "-")):
                    key = key[1:]
            # Use real field name
            key = QuerySet._translate_field_name(self._document, key)
            index_list.append((key, direction))
        self._collection.ensure_index(index_list)
        return self

    def __call__(self, *q_objs, **query):
        """Filter the selected documents by calling the 
        :class:`~mongoengine.QuerySet` with a query.

        :param q_objs: :class:`~mongoengine.Q` objects to be used in the query
        :param query: Django-style query keyword arguments
        """
        for q in q_objs:
            self._where_clauses.append(q.as_js(self._document))
        query = QuerySet._transform_query(_doc_cls=self._document, **query)
        self._query.update(query)
        return self
        
    def filter(self, *q_objs, **query):
        return self.__call__(**query)

    @property
    def _collection(self):
        """Property that returns the collection object. This allows us to
        perform operations only if the collection is accessed.
        """
        if not self._accessed_collection:
            self._accessed_collection = True

            # Ensure document-defined indexes are created
            if self._document._meta['indexes']:
                for key_or_list in self._document._meta['indexes']:
                    self.ensure_index(key_or_list)

            # Ensure indexes created by uniqueness constraints
            for index in self._document._meta['unique_indexes']:
                self._collection.ensure_index(index, unique=True)

            # If _types is being used (for polymorphism), it needs an index
            if '_types' in self._query:
                self._collection.ensure_index('_types')
        return self._collection_obj

    @property
    def _cursor(self):
        if not self._cursor_obj:
            self._cursor_obj = self._collection.find(self._query)
            # Apply where clauses to cursor
            for js in self._where_clauses:
                self._cursor_obj.where(js)
{'fraction_non_alphanumeric': 0.1209410438825104, 'fraction_numerical': 0.0021131224906670424, 'mean_word_length': 1.9566847147022075, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, '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': '5289339', 'n_tokens_mistral': 4235, 'n_tokens_neox': 4049, 'n_words': 1097}