From 628a8dacfaeea98e94882620b334f097a0911fb1 Mon Sep 17 00:00:00 2001 From: Thomas Steinacher Date: Sun, 16 Jun 2013 23:28:31 -0700 Subject: [PATCH] Fix select_related for documents --- mongoengine/dereference.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mongoengine/dereference.py b/mongoengine/dereference.py index b9d79e6..79f755f 100644 --- a/mongoengine/dereference.py +++ b/mongoengine/dereference.py @@ -89,6 +89,8 @@ class DeReference(object): v = getattr(item, field_name) if isinstance(v, (DBRef)): reference_map.setdefault(field.document_type, []).append(v.id) + elif isinstance(v, Document) and getattr(v, '_lazy', False): + reference_map.setdefault(field.document_type, []).append(v.pk) 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: @@ -196,6 +198,8 @@ class DeReference(object): v = data[k]._internal_data.get(field_name, None) if isinstance(v, (DBRef)): data[k]._internal_data[field_name] = self.object_map.get(v.id, v) + elif isinstance(v, Document) and getattr(v, '_lazy', False): + data[k]._internal_data[field_name] = self.object_map.get(v.pk, v) elif isinstance(v, (dict, SON)) and '_ref' in v: data[k]._internal_data[field_name] = self.object_map.get(v['_ref'].id, v) elif isinstance(v, dict) and depth <= self.max_depth: