Updated .only() behaviour - now like exclude it is chainable (#202)

mongoengine/queryset/field_list.py

@@ -7,11 +7,20 @@ class QueryFieldList(object):
     ONLY = 1
     EXCLUDE = 0
 
-    def __init__(self, fields=[], value=ONLY, always_include=[]):
+    def __init__(self, fields=None, value=ONLY, always_include=None, _only_called=False):
+        """The QueryFieldList builder
+
+        :param fields: A list of fields used in `.only()` or `.exclude()`
+        :param value: How to handle the fields; either `ONLY` or `EXCLUDE`
+        :param always_include: Any fields to always_include eg `_cls`
+        :param _only_called: Has `.only()` been called?  If so its a set of fields
+           otherwise it performs a union.
+        """
         self.value = value
-        self.fields = set(fields)
-        self.always_include = set(always_include)
+        self.fields = set(fields or [])
+        self.always_include = set(always_include or [])
         self._id = None
+        self._only_called = _only_called
         self.slice = {}
 
     def __add__(self, f):
@@ -26,7 +35,10 @@ class QueryFieldList(object):
             self.slice = {}
         elif self.value is self.ONLY and f.value is self.ONLY:
             self._clean_slice()
-            self.fields = self.fields.intersection(f.fields)
+            if self._only_called:
+                self.fields = self.fields.union(f.fields)
+            else:
+                self.fields = f.fields
         elif self.value is self.EXCLUDE and f.value is self.EXCLUDE:
             self.fields = self.fields.union(f.fields)
             self._clean_slice()
@@ -46,6 +58,9 @@ class QueryFieldList(object):
                 self.fields = self.fields.union(self.always_include)
             else:
                 self.fields -= self.always_include
+
+        if getattr(f, '_only_called', False):
+            self._only_called = True
         return self
 
     def __nonzero__(self):

mongoengine/queryset/queryset.py

@@ -624,19 +624,35 @@ class QuerySet(object):
 
             post = BlogPost.objects(...).only("title", "author.name")
 
+        .. note :: `only()` is chainable and will perform a union ::
+            So with the following it will fetch both: `title` and `author.name`::
+
+                post = BlogPost.objects.only("title").only("author.name")
+
+        :func:`~mongoengine.queryset.QuerySet.all_fields` will reset any
+        field filters.
+
         :param fields: fields to include
 
         .. versionadded:: 0.3
         .. versionchanged:: 0.5 - Added subfield support
         """
         fields = dict([(f, QueryFieldList.ONLY) for f in fields])
-        return self.fields(**fields)
+        return self.fields(True, **fields)
 
     def exclude(self, *fields):
         """Opposite to .only(), exclude some document's fields. ::
 
             post = BlogPost.objects(...).exclude("comments")
 
+        .. note :: `exclude()` is chainable and will perform a union ::
+            So with the following it will exclude both: `title` and `author.name`::
+
+                post = BlogPost.objects.exclude("title").exclude("author.name")
+
+        :func:`~mongoengine.queryset.QuerySet.all_fields` will reset any
+        field filters.
+
         :param fields: fields to exclude
 
         .. versionadded:: 0.5
@@ -644,7 +660,7 @@ class QuerySet(object):
         fields = dict([(f, QueryFieldList.EXCLUDE) for f in fields])
         return self.fields(**fields)
 
-    def fields(self, **kwargs):
+    def fields(self, _only_called=False, **kwargs):
         """Manipulate how you load this document's fields.  Used by `.only()`
         and `.exclude()` to manipulate which fields to retrieve.  Fields also
         allows for a greater level of control for example:
@@ -678,7 +694,8 @@ class QuerySet(object):
         for value, group in itertools.groupby(fields, lambda x: x[1]):
             fields = [field for field, value in group]
             fields = queryset._fields_to_dbfields(fields)
-            queryset._loaded_fields += QueryFieldList(fields, value=value)
+            queryset._loaded_fields += QueryFieldList(fields, value=value, _only_called=_only_called)
+
         return queryset
 
     def all_fields(self):