11150999376

Committed 02 Oct 2024 08:19PM UTC coverage: 94.447%. First build

Build # 11150999376

Build Type

push

github

Committed by

bagerard

Commit Message

refactor _document_registry + log a warning when user register multiple Document classes with the same name (only flagging when this happens in different module)

Run Details

42 of 44 new or added lines in 6 files covered. (95.45%)

5341 of 5655 relevant lines covered (94.45%)

1.88 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

96.65

/mongoengine/dereference.py

from bson import SON, DBRef

from mongoengine.base import (
    BaseDict,
    BaseList,
    EmbeddedDocumentList,
    TopLevelDocumentMetaclass,
    _DocumentRegistry,
)
from mongoengine.base.datastructures import LazyReference
from mongoengine.connection import _get_session, get_db
from mongoengine.document import Document, EmbeddedDocument
from mongoengine.fields import (
    DictField,
    ListField,
    MapField,
    ReferenceField,
)
from mongoengine.queryset import QuerySet


class DeReference:
    def __call__(self, items, max_depth=1, instance=None, name=None):
        """
        Cheaply dereferences the items to a set depth.
        Also handles the conversion 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, str):
            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, EmbeddedDocument, 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:
                    # We must turn the ObjectIds into DBRefs

                    # Recursively dig into the sub items of a list/dict
                    # to turn the ObjectIds into DBRefs
                    def _get_items_from_list(items):
                        new_items = []
                        for v in items:
                            value = v
                            if isinstance(v, dict):
                                value = _get_items_from_dict(v)
                            elif isinstance(v, list):
                                value = _get_items_from_list(v)
                            elif not isinstance(v, (DBRef, Document)):
                                value = field.to_python(v)
                            new_items.append(value)
                        return new_items

                    def _get_items_from_dict(items):
                        new_items = {}
                        for k, v in items.items():
                            value = v
                            if isinstance(v, list):
                                value = _get_items_from_list(v)
                            elif isinstance(v, dict):
                                value = _get_items_from_dict(v)
                            elif not isinstance(v, (DBRef, Document)):
                                value = field.to_python(v)
                            new_items[k] = value
                        return new_items

                    if not hasattr(items, "items"):
                        items = _get_items_from_list(items)
                    else:
                        items = _get_items_from_dict(items)

        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 isinstance(items, dict):
            iterator = items.values()
        else:
            iterator = items

        # Recursively find dbreferences
        depth += 1
        for item in iterator:
            if isinstance(item, (Document, EmbeddedDocument)):
                for field_name, field in item._fields.items():
                    v = item._data.get(field_name, None)
                    if isinstance(v, LazyReference):
                        # LazyReference inherits DBRef but should not be dereferenced here !
                        continue
                    elif isinstance(v, DBRef):
                        reference_map.setdefault(field.document_type, set()).add(v.id)
                    elif isinstance(v, (dict, SON)) and "_ref" in v:
                        reference_map.setdefault(
                            _DocumentRegistry.get(v["_cls"]), set()
                        ).add(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.items():
                            if isinstance(
                                field_cls, (Document, TopLevelDocumentMetaclass)
                            ):
                                key = field_cls
                            reference_map.setdefault(key, set()).update(refs)
            elif isinstance(item, LazyReference):
                # LazyReference inherits DBRef but should not be dereferenced here !
                continue
            elif isinstance(item, DBRef):
                reference_map.setdefault(item.collection, set()).add(item.id)
            elif isinstance(item, (dict, SON)) and "_ref" in item:
                reference_map.setdefault(
                    _DocumentRegistry.get(item["_cls"]), set()
                ).add(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.items():
                    reference_map.setdefault(key, set()).update(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.items():
            # we use getattr instead of hasattr because hasattr swallows any exception under python2
            # so it could hide nasty things without raising exceptions (cfr bug #1688))
            ref_document_cls_exists = getattr(collection, "objects", None) is not None

            if ref_document_cls_exists:
                col_name = collection._get_collection_name()
                refs = [
                    dbref for dbref in dbrefs if (col_name, dbref) not in object_map
                ]
                references = collection.objects.in_bulk(refs)
                for key, doc in references.items():
                    object_map[(col_name, key)] = doc
            else:  # Generic reference: use the refs data to convert to document
                if isinstance(doc_type, (ListField, DictField, MapField)):
                    continue

                refs = [
                    dbref for dbref in dbrefs if (collection, dbref) not in object_map
                ]

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

    def _attach_objects(self, items, depth=0, instance=None, name=None):
        """
        Recursively finds all db references to be dereferenced

        :param items: The iterable (dict, list, queryset)
        :param depth: The current depth of recursion
        :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`
        """
        if not items:
            if isinstance(items, (BaseDict, BaseList)):
                return items

            if instance:
                if isinstance(items, dict):
                    return BaseDict(items, instance, name)
                else:
                    return BaseList(items, instance, name)

        if isinstance(items, (dict, SON)):
            if "_ref" in items:
                return self.object_map.get(
                    (items["_ref"].collection, items["_ref"].id), items
                )
            elif "_cls" in items:
                doc = _DocumentRegistry.get(items["_cls"])._from_son(items)
                _cls = doc._data.pop("_cls", None)
                del items["_cls"]
                doc._data = self._attach_objects(doc._data, depth, doc, None)
                if _cls is not None:
                    doc._data["_cls"] = _cls
                return doc

        if not hasattr(items, "items"):
            is_list = True
            list_type = BaseList
            if isinstance(items, EmbeddedDocumentList):
                list_type = EmbeddedDocumentList
            as_tuple = isinstance(items, tuple)
            iterator = enumerate(items)
            data = []
        else:
            is_list = False
            iterator = items.items()
            data = {}

        depth += 1
        for k, v in iterator:
            if is_list:
                data.append(v)
            else:
                data[k] = v

            if k in self.object_map and not is_list:
                data[k] = self.object_map[k]
            elif isinstance(v, (Document, EmbeddedDocument)):
                for field_name in v._fields:
                    v = data[k]._data.get(field_name, None)
                    if isinstance(v, DBRef):
                        data[k]._data[field_name] = self.object_map.get(
                            (v.collection, v.id), v
                        )
                    elif isinstance(v, (dict, SON)) and "_ref" in v:
                        data[k]._data[field_name] = self.object_map.get(
                            (v["_ref"].collection, v["_ref"].id), v
                        )
                    elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:
                        item_name = f"{name}.{k}.{field_name}"
                        data[k]._data[field_name] = self._attach_objects(
                            v, depth, instance=instance, name=item_name
                        )
            elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:
                item_name = f"{name}.{k}" if name else name
                data[k] = self._attach_objects(
                    v, depth - 1, instance=instance, name=item_name
                )
            elif isinstance(v, DBRef) and hasattr(v, "id"):
                data[k] = self.object_map.get((v.collection, v.id), v)

        if instance and name:
            if is_list:
                return tuple(data) if as_tuple else list_type(data, instance, name)
            return BaseDict(data, instance, name)
        depth += 1
        return data

1	from bson import SON, DBRef	2✔
2
3	from mongoengine.base import (	2✔
4	BaseDict,
5	BaseList,
6	EmbeddedDocumentList,
7	TopLevelDocumentMetaclass,
8	_DocumentRegistry,
9	)
10	from mongoengine.base.datastructures import LazyReference	2✔
11	from mongoengine.connection import _get_session, get_db	2✔
12	from mongoengine.document import Document, EmbeddedDocument	2✔
13	from mongoengine.fields import (	2✔
14	DictField,
15	ListField,
16	MapField,
17	ReferenceField,
18	)
19	from mongoengine.queryset import QuerySet	2✔
20
21
22	class DeReference:	2✔
23	def __call__(self, items, max_depth=1, instance=None, name=None):	2✔
24	"""
25	Cheaply dereferences the items to a set depth.
26	Also handles the conversion of complex data types.
27
28	:param items: The iterable (dict, list, queryset) to be dereferenced.
29	:param max_depth: The maximum depth to recurse to
30	:param instance: The owning instance used for tracking changes by
31	:class:`~mongoengine.base.ComplexBaseField`
32	:param name: The name of the field, used for tracking changes by
33	:class:`~mongoengine.base.ComplexBaseField`
34	:param get: A boolean determining if being called by __get__
35	"""
36	if items is None or isinstance(items, str):	2✔
37	return items	×
38
39	# cheapest way to convert a queryset to a list
40	# list(queryset) uses a count() query to determine length
41	if isinstance(items, QuerySet):	2✔
42	items = [i for i in items]	2✔
43
44	self.max_depth = max_depth	2✔
45	doc_type = None	2✔
46
47	if instance and isinstance(	2✔
48	instance, (Document, EmbeddedDocument, TopLevelDocumentMetaclass)
49	):
50	doc_type = instance._fields.get(name)	2✔
51	while hasattr(doc_type, "field"):	2✔
52	doc_type = doc_type.field	2✔
53
54	if isinstance(doc_type, ReferenceField):	2✔
55	field = doc_type	2✔
56	doc_type = doc_type.document_type	2✔
57	is_list = not hasattr(items, "items")	2✔
58
59	if is_list and all(i.__class__ == doc_type for i in items):	2✔
60	return items	2✔
61	elif not is_list and all(	2✔
62	i.__class__ == doc_type for i in items.values()
63	):
64	return items	2✔
65	elif not field.dbref:	2✔
66	# We must turn the ObjectIds into DBRefs
67
68	# Recursively dig into the sub items of a list/dict
69	# to turn the ObjectIds into DBRefs
70	def _get_items_from_list(items):	2✔
71	new_items = []	2✔
72	for v in items:	2✔
73	value = v	2✔
74	if isinstance(v, dict):	2✔
75	value = _get_items_from_dict(v)	2✔
76	elif isinstance(v, list):	2✔
77	value = _get_items_from_list(v)	2✔
78	elif not isinstance(v, (DBRef, Document)):	2✔
79	value = field.to_python(v)	2✔
80	new_items.append(value)	2✔
81	return new_items	2✔
82
83	def _get_items_from_dict(items):	2✔
84	new_items = {}	2✔
85	for k, v in items.items():	2✔
86	value = v	2✔
87	if isinstance(v, list):	2✔
88	value = _get_items_from_list(v)	2✔
89	elif isinstance(v, dict):	2✔
90	value = _get_items_from_dict(v)	2✔
91	elif not isinstance(v, (DBRef, Document)):	2✔
92	value = field.to_python(v)	×
93	new_items[k] = value	2✔
94	return new_items	2✔
95
96	if not hasattr(items, "items"):	2✔
97	items = _get_items_from_list(items)	2✔
98	else:
99	items = _get_items_from_dict(items)	2✔
100
101	self.reference_map = self._find_references(items)	2✔
102	self.object_map = self._fetch_objects(doc_type=doc_type)	2✔
103	return self._attach_objects(items, 0, instance, name)	2✔
104
105	def _find_references(self, items, depth=0):	2✔
106	"""
107	Recursively finds all db references to be dereferenced
108
109	:param items: The iterable (dict, list, queryset)
110	:param depth: The current depth of recursion
111	"""
112	reference_map = {}	2✔
113	if not items or depth >= self.max_depth:	2✔
114	return reference_map	2✔
115
116	# Determine the iterator to use
117	if isinstance(items, dict):	2✔
118	iterator = items.values()	2✔
119	else:
120	iterator = items	2✔
121
122	# Recursively find dbreferences
123	depth += 1	2✔
124	for item in iterator:	2✔
125	if isinstance(item, (Document, EmbeddedDocument)):	2✔
126	for field_name, field in item._fields.items():	2✔
127	v = item._data.get(field_name, None)	2✔
128	if isinstance(v, LazyReference):	2✔
129	# LazyReference inherits DBRef but should not be dereferenced here !
130	continue	2✔
131	elif isinstance(v, DBRef):	2✔
132	reference_map.setdefault(field.document_type, set()).add(v.id)	2✔
133	elif isinstance(v, (dict, SON)) and "_ref" in v:	2✔
134	reference_map.setdefault(	2✔
135	_DocumentRegistry.get(v["_cls"]), set()
136	).add(v["_ref"].id)
137	elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:	2✔
138	field_cls = getattr(	2✔
139	getattr(field, "field", None), "document_type", None
140	)
141	references = self._find_references(v, depth)	2✔
142	for key, refs in references.items():	2✔
143	if isinstance(	2✔
144	field_cls, (Document, TopLevelDocumentMetaclass)
145	):
146	key = field_cls	2✔
147	reference_map.setdefault(key, set()).update(refs)	2✔
148	elif isinstance(item, LazyReference):	2✔
149	# LazyReference inherits DBRef but should not be dereferenced here !
150	continue	2✔
151	elif isinstance(item, DBRef):	2✔
152	reference_map.setdefault(item.collection, set()).add(item.id)	2✔
153	elif isinstance(item, (dict, SON)) and "_ref" in item:	2✔
154	reference_map.setdefault(	2✔
155	_DocumentRegistry.get(item["_cls"]), set()
156	).add(item["_ref"].id)
157	elif isinstance(item, (dict, list, tuple)) and depth - 1 <= self.max_depth:	2✔
158	references = self._find_references(item, depth - 1)	2✔
159	for key, refs in references.items():	2✔
160	reference_map.setdefault(key, set()).update(refs)	2✔
161
162	return reference_map	2✔
163
164	def _fetch_objects(self, doc_type=None):	2✔
165	"""Fetch all references and convert to their document objects"""
166	object_map = {}	2✔
167	for collection, dbrefs in self.reference_map.items():	2✔
168	# we use getattr instead of hasattr because hasattr swallows any exception under python2
169	# so it could hide nasty things without raising exceptions (cfr bug #1688))
170	ref_document_cls_exists = getattr(collection, "objects", None) is not None	2✔
171
172	if ref_document_cls_exists:	2✔
173	col_name = collection._get_collection_name()	2✔
174	refs = [	2✔
175	dbref for dbref in dbrefs if (col_name, dbref) not in object_map
176	]
177	references = collection.objects.in_bulk(refs)	2✔
178	for key, doc in references.items():	2✔
179	object_map[(col_name, key)] = doc	2✔
180	else: # Generic reference: use the refs data to convert to document
181	if isinstance(doc_type, (ListField, DictField, MapField)):	2✔
182	continue	×
183
184	refs = [	2✔
185	dbref for dbref in dbrefs if (collection, dbref) not in object_map
186	]
187
188	if doc_type:	2✔
189	references = doc_type._get_db()[collection].find(	2✔
190	{"_id": {"$in": refs}}, session=_get_session()
191	)
192	for ref in references:	2✔
193	doc = doc_type._from_son(ref)	2✔
194	object_map[(collection, doc.id)] = doc	2✔
195	else:
196	references = get_db()[collection].find(	2✔
197	{"_id": {"$in": refs}}, session=_get_session()
198	)
199	for ref in references:	2✔
200	if "_cls" in ref:	2✔
NEW 201	doc = _DocumentRegistry.get(ref["_cls"])._from_son(ref)	×
202	elif doc_type is None:	2✔
203	doc = _DocumentRegistry.get(	2✔
204	"".join(x.capitalize() for x in collection.split("_"))
205	)._from_son(ref)
206	else:
207	doc = doc_type._from_son(ref)	×
208	object_map[(collection, doc.id)] = doc	2✔
209	return object_map	2✔
210
211	def _attach_objects(self, items, depth=0, instance=None, name=None):	2✔
212	"""
213	Recursively finds all db references to be dereferenced
214
215	:param items: The iterable (dict, list, queryset)
216	:param depth: The current depth of recursion
217	:param instance: The owning instance used for tracking changes by
218	:class:`~mongoengine.base.ComplexBaseField`
219	:param name: The name of the field, used for tracking changes by
220	:class:`~mongoengine.base.ComplexBaseField`
221	"""
222	if not items:	2✔
223	if isinstance(items, (BaseDict, BaseList)):	2✔
224	return items	2✔
225
226	if instance:	2✔
227	if isinstance(items, dict):	2✔
228	return BaseDict(items, instance, name)	2✔
229	else:
230	return BaseList(items, instance, name)	2✔
231
232	if isinstance(items, (dict, SON)):	2✔
233	if "_ref" in items:	2✔
234	return self.object_map.get(	2✔
235	(items["_ref"].collection, items["_ref"].id), items
236	)
237	elif "_cls" in items:	2✔
238	doc = _DocumentRegistry.get(items["_cls"])._from_son(items)	2✔
239	_cls = doc._data.pop("_cls", None)	2✔
240	del items["_cls"]	2✔
241	doc._data = self._attach_objects(doc._data, depth, doc, None)	2✔
242	if _cls is not None:	2✔
243	doc._data["_cls"] = _cls	2✔
244	return doc	2✔
245
246	if not hasattr(items, "items"):	2✔
247	is_list = True	2✔
248	list_type = BaseList	2✔
249	if isinstance(items, EmbeddedDocumentList):	2✔
250	list_type = EmbeddedDocumentList	2✔
251	as_tuple = isinstance(items, tuple)	2✔
252	iterator = enumerate(items)	2✔
253	data = []	2✔
254	else:
255	is_list = False	2✔
256	iterator = items.items()	2✔
257	data = {}	2✔
258
259	depth += 1	2✔
260	for k, v in iterator:	2✔
261	if is_list:	2✔
262	data.append(v)	2✔
263	else:
264	data[k] = v	2✔
265
266	if k in self.object_map and not is_list:	2✔
267	data[k] = self.object_map[k]	×
268	elif isinstance(v, (Document, EmbeddedDocument)):	2✔
269	for field_name in v._fields:	2✔
270	v = data[k]._data.get(field_name, None)	2✔
271	if isinstance(v, DBRef):	2✔
272	data[k]._data[field_name] = self.object_map.get(	2✔
273	(v.collection, v.id), v
274	)
275	elif isinstance(v, (dict, SON)) and "_ref" in v:	2✔
276	data[k]._data[field_name] = self.object_map.get(	2✔
277	(v["_ref"].collection, v["_ref"].id), v
278	)
279	elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:	2✔
280	item_name = f"{name}.{k}.{field_name}"	2✔
281	data[k]._data[field_name] = self._attach_objects(	2✔
282	v, depth, instance=instance, name=item_name
283	)
284	elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:	2✔
285	item_name = f"{name}.{k}" if name else name	2✔
286	data[k] = self._attach_objects(	2✔
287	v, depth - 1, instance=instance, name=item_name
288	)
289	elif isinstance(v, DBRef) and hasattr(v, "id"):	2✔
290	data[k] = self.object_map.get((v.collection, v.id), v)	2✔
291
292	if instance and name:	2✔
293	if is_list:	2✔
294	return tuple(data) if as_tuple else list_type(data, instance, name)	2✔
295	return BaseDict(data, instance, name)	2✔
296	depth += 1	2✔
297	return data	2✔

bagerard / mongoengine / 11150999376

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous