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

This commit is contained in:
Ross Lawley
2013-04-23 14:06:29 +00:00
parent 81c7007f80
commit e2f3406e89
5 changed files with 73 additions and 25 deletions

View File

@@ -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):