<commit_before>from bson import DBRef, SON

from base import (BaseDict, BaseList, TopLevelDocumentMetaclass, get_document)
from fields import (ReferenceField, ListField, DictField, MapField)
from connection import get_db
from queryset import QuerySet
from document import Document


class DeReference(object):

    def __call__(self, items, max_depth=1, instance=None, name=None):
        """
        Cheaply dereferences the items to a set depth.
        Also handles the convertion of complex data types.

        :param items: The iterable (dict, list, queryset) to be dereferenced.
        :param max_depth: The maximum depth to recurse to
        :param instance: The owning instance used for tracking changes by
            :class:`~mongoengine.base.ComplexBaseField`
        :param name: The name of the field, used for tracking changes by
            :class:`~mongoengine.base.ComplexBaseField`
        :param get: A boolean determining if being called by __get__
        """
        if items is None or isinstance(items, basestring):
            return items

        # cheapest way to convert a queryset to a list
        # list(queryset) uses a count() query to determine length
        if isinstance(items, QuerySet):
            items = [i for i in items]

        self.max_depth = max_depth
        doc_type = None

        if instance and isinstance(instance, (Document, TopLevelDocumentMetaclass)):
            doc_type = instance._fields.get(name)
            if hasattr(doc_type, 'field'):
                doc_type = doc_type.field

            if isinstance(doc_type, ReferenceField):
                field = doc_type
                doc_type = doc_type.document_type
                is_list = not hasattr(items, 'items')

                if is_list and all([i.__class__ == doc_type for i in items]):
                    return items
                elif not is_list and all([i.__class__ == doc_type
                                         for i in items.values()]):
                    return items
                elif not field.dbref:
                    if not hasattr(items, 'items'):
                        items = [field.to_python(v)
                             if not isinstance(v, (DBRef, Document)) else v
                             for v in items]
                    else:
                        items = dict([
                            (k, field.to_python(v))
                            if not isinstance(v, (DBRef, Document)) else (k, v)
                            for k, v in items.iteritems()]
                        )

        self.reference_map = self._find_references(items)
        self.object_map = self._fetch_objects(doc_type=doc_type)
        return self._attach_objects(items, 0, instance, name)

    def _find_references(self, items, depth=0):
        """
        Recursively finds all db references to be dereferenced

        :param items: The iterable (dict, list, queryset)
        :param depth: The current depth of recursion
        """
        reference_map = {}
        if not items or depth >= self.max_depth:
            return reference_map

        # Determine the iterator to use
        if not hasattr(items, 'items'):
            iterator = enumerate(items)
        else:
            iterator = items.iteritems()

        # Recursively find dbreferences
        depth += 1
        for k, item in iterator:
            if isinstance(item, Document):
                for field_name, field in item._fields.iteritems():
                    v = item._data.get(field_name, None)
                    if isinstance(v, (DBRef)):
                        reference_map.setdefault(field.document_type, []).append(v.id)
                    elif isinstance(v, (dict, SON)) and '_ref' in v:
                        reference_map.setdefault(get_document(v['_cls']), []).append(v['_ref'].id)
                    elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:
                        field_cls = getattr(getattr(field, 'field', None), 'document_type', None)
                        references = self._find_references(v, depth)
                        for key, refs in references.iteritems():
                            if isinstance(field_cls, (Document, TopLevelDocumentMetaclass)):
                                key = field_cls
                            reference_map.setdefault(key, []).extend(refs)
            elif isinstance(item, (DBRef)):
                reference_map.setdefault(item.collection, []).append(item.id)
            elif isinstance(item, (dict, SON)) and '_ref' in item:
                reference_map.setdefault(get_document(item['_cls']), []).append(item['_ref'].id)
            elif isinstance(item, (dict, list, tuple)) and depth - 1 <= self.max_depth:
                references = self._find_references(item, depth - 1)
                for key, refs in references.iteritems():
                    reference_map.setdefault(key, []).extend(refs)

        return reference_map

    def _fetch_objects(self, doc_type=None):
        """Fetch all references and convert to their document objects
        """
        object_map = {}
        for col, dbrefs in self.reference_map.iteritems():
            keys = object_map.keys()
            refs = list(set([dbref for dbref in dbrefs if unicode(dbref).encode('utf-8') not in keys]))
            if hasattr(col, 'objects'):  # We have a document class for the refs
                references = col.objects.in_bulk(refs)
                for key, doc in references.iteritems():
                    object_map[key] = doc
            else:  # Generic reference: use the refs data to convert to document
                if isinstance(doc_type, (ListField, DictField, MapField,)):
                    continue

                if doc_type:
                    references = doc_type._get_db()[col].find({'_id': {'$in': refs}})
                    for ref in references:
                        doc = doc_type._from_son(ref)
                        object_map[doc.id] = doc
                else:
                    references = get_db()[col].find({'_id': {'$in': refs}})
                    for ref in references:
                        if '_cls' in ref:
                            doc = get_document(ref["_cls"])._from_son(ref)
                        elif doc_type is None:
                            doc = get_document(
                                ''.join(x.capitalize()
                                    for x in col.split('_')))._from_son(ref)
                        else:
                            doc = doc_type._from_son(ref)
                        object_map[doc.id] = doc
        return object_map
<commit_msg>potential fix for dereferencing nested lists<commit_after>from bson import DBRef, SON

from base import (BaseDict, BaseList, TopLevelDocumentMetaclass, get_document)
from fields import (ReferenceField, ListField, DictField, MapField)
from connection import get_db
from queryset import QuerySet
from document import Document


class DeReference(object):

    def __call__(self, items, max_depth=1, instance=None, name=None):
        """
        Cheaply dereferences the items to a set depth.
        Also handles the convertion of complex data types.

        :param items: The iterable (dict, list, queryset) to be dereferenced.
        :param max_depth: The maximum depth to recurse to
        :param instance: The owning instance used for tracking changes by
            :class:`~mongoengine.base.ComplexBaseField`
        :param name: The name of the field, used for tracking changes by
            :class:`~mongoengine.base.ComplexBaseField`
        :param get: A boolean determining if being called by __get__
        """
        if items is None or isinstance(items, basestring):
            return items

        # cheapest way to convert a queryset to a list
        # list(queryset) uses a count() query to determine length
        if isinstance(items, QuerySet):
            items = [i for i in items]

        self.max_depth = max_depth
        doc_type = None

        if instance and isinstance(instance, (Document, TopLevelDocumentMetaclass)):
            doc_type = instance._fields.get(name)
            while hasattr(doc_type, 'field'):
                doc_type = doc_type.field

            if isinstance(doc_type, ReferenceField):
                field = doc_type
                doc_type = doc_type.document_type
                is_list = not hasattr(items, 'items')

                if is_list and all([i.__class__ == doc_type for i in items]):
                    return items
                elif not is_list and all([i.__class__ == doc_type
                                         for i in items.values()]):
                    return items
                elif not field.dbref:
                    if not hasattr(items, 'items'):

                        def _get_items(items):
                            new_items = []
                            for v in items:
                                if isinstance(v, list):
                                    new_items.append(_get_items(v))
                                elif not isinstance(v, (DBRef, Document)):
                                    new_items.append(field.to_python(v))
                                else:
                                    new_items.append(v)
                            return new_items

                        items = _get_items(items)
                    else:
                        items = dict([
                            (k, field.to_python(v))
                            if not isinstance(v, (DBRef, Document)) else (k, v)
                            for k, v in items.iteritems()]
                        )

        self.reference_map = self._find_references(items)
        self.object_map = self._fetch_objects(doc_type=doc_type)
        return self._attach_objects(items, 0, instance, name)

    def _find_references(self, items, depth=0):
        """
        Recursively finds all db references to be dereferenced

        :param items: The iterable (dict, list, queryset)
        :param depth: The current depth of recursion
        """
        reference_map = {}
        if not items or depth >= self.max_depth:
            return reference_map

        # Determine the iterator to use
        if not hasattr(items, 'items'):
            iterator = enumerate(items)
        else:
            iterator = items.iteritems()

        # Recursively find dbreferences
        depth += 1
        for k, item in iterator:
            if isinstance(item, Document):
                for field_name, field in item._fields.iteritems():
                    v = item._data.get(field_name, None)
                    if isinstance(v, (DBRef)):
                        reference_map.setdefault(field.document_type, []).append(v.id)
                    elif isinstance(v, (dict, SON)) and '_ref' in v:
                        reference_map.setdefault(get_document(v['_cls']), []).append(v['_ref'].id)
                    elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:
                        field_cls = getattr(getattr(field, 'field', None), 'document_type', None)
                        references = self._find_references(v, depth)
                        for key, refs in references.iteritems():
                            if isinstance(field_cls, (Document, TopLevelDocumentMetaclass)):
                                key = field_cls
                            reference_map.setdefault(key, []).extend(refs)
            elif isinstance(item, (DBRef)):
                reference_map.setdefault(item.collection, []).append(item.id)
            elif isinstance(item, (dict, SON)) and '_ref' in item:
                reference_map.setdefault(get_document(item['_cls']), []).append(item['_ref'].id)
            elif isinstance(item, (dict, list, tuple)) and depth - 1 <= self.max_depth:
                references = self._find_references(item, depth - 1)
                for key, refs in references.iteritems():
                    reference_map.setdefault(key, []).extend(refs)

        return reference_map

    def _fetch_objects(self, doc_type=None):
        """Fetch all references and convert to their document objects
        """
        object_map = {}
        for collection, dbrefs in self.reference_map.iteritems():
            keys = object_map.keys()
            refs = list(set([dbref for dbref in dbrefs if unicode(dbref).encode('utf-8') not in keys]))
            if hasattr(collection, 'objects'):  # We have a document class for the refs
                references = collection.objects.in_bulk(refs)
                for key, doc in references.iteritems():
                    object_map[key] = doc
            else:  # Generic reference: use the refs data to convert to document
                if isinstance(doc_type, (ListField, DictField, MapField,)):
                    continue

                if doc_type:
                    references = doc_type._get_db()[collection].find({'_id': {'$in': refs}})
                    for ref in references:
                        doc = doc_type._from_son(ref)
                        object_map[doc.id] = doc
                else:
                    references = get_db()[collection].find({'_id': {'$in': refs}})
                    for ref in references:
                        if '_cls' in ref:
                            doc = get_document(ref["_cls"])._from_son(ref)
                        elif doc_type is None:
                            doc = get_document(
                                ''.join(x.capitalize()
                                    for x in collection.split('_')))._from_son(ref)
                        else:
                            doc = doc_type._from_son(ref)
                        object_map[doc.id] = doc
        return object_map

{'fraction_non_alphanumeric': 0.08854542828253648, 'fraction_numerical': 0.0010111223458038423, 'mean_word_length': 1.6033088926489942, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7289547', 'n_tokens_mistral': 3589, 'n_tokens_neox': 3419, 'n_words': 1029}