From ed5fba6b0f8794c5d0c2b66a4ede2fdd6346ec75 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wilson=20J=C3=BAnior?= <wilsonpjunior@gmail.com>
Date: Tue, 13 Dec 2011 07:46:49 -0200
Subject: [PATCH] support for embedded fields

---
 mongoengine/queryset.py | 14 ++++++++++++--
 tests/queryset.py       | 41 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index 1b3257ac..1a2ff1c3 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -320,12 +320,22 @@ class SelectResult(object):
     """
     def __init__(self, cursor, fields):
         self._cursor = cursor
-        self._fields = fields
+        self._fields = [f.split('.') for f in fields]
+        
+    def _get_value(self, keys, data):
+        for key in keys:
+            if data:
+                data = data.get(key)
+            else:
+                break
+
+        return data
 
     def next(self):
         try:
             data = self._cursor.next()
-            return [data.get(f) for f in self._fields]
+            return [self._get_value(f, data)
+                    for f in self._fields]
         except StopIteration, e:
             self.rewind()
             raise e
diff --git a/tests/queryset.py b/tests/queryset.py
index 0f5d84cf..991fb00a 100644
--- a/tests/queryset.py
+++ b/tests/queryset.py
@@ -2955,5 +2955,46 @@ class QueryFieldListTest(unittest.TestCase):
                 [u'Eliana'],
                 [u'Wilson']])
 
+    def test_select_embedded(self):
+        class Profile(EmbeddedDocument):
+            name = StringField()
+            age = IntField()
+
+        class Locale(EmbeddedDocument):
+            city = StringField()
+            country = StringField()
+
+        class Person(Document):
+            profile = EmbeddedDocumentField(Profile)
+            locale = EmbeddedDocumentField(Locale)
+
+        Person.drop_collection()
+
+        Person(profile=Profile(name="Wilson Jr", age=19),
+               locale=Locale(city="Corumba-GO", country="Brazil")).save()
+
+        Person(profile=Profile(name="Gabriel Falcao", age=23),
+               locale=Locale(city="New York", country="USA")).save()
+
+        Person(profile=Profile(name="Lincoln de souza", age=28),
+               locale=Locale(city="Belo Horizonte", country="Brazil")).save()
+
+        Person(profile=Profile(name="Walter cruz", age=30),
+               locale=Locale(city="Brasilia", country="Brazil")).save()
+
+        self.assertEqual(
+            list(Person.objects.order_by('profile.age').select('profile.name')),
+            [[u'Wilson Jr'], [u'Gabriel Falcao'],
+             [u'Lincoln de souza'], [u'Walter cruz']])
+
+        ulist = list(Person.objects.order_by('locale.city')
+                     .select('profile.name', 'profile.age', 'locale.city'))
+        self.assertEqual(ulist,
+                         [[u'Lincoln de souza', 28, u'Belo Horizonte'],
+                          [u'Walter cruz', 30, u'Brasilia'],
+                          [u'Wilson Jr', 19, u'Corumba-GO'],
+                          [u'Gabriel Falcao', 23, u'New York']])
+            
+
 if __name__ == '__main__':
     unittest.main()