Skip to content

Commit

Permalink
.iterator() should respect return format
Browse files Browse the repository at this point in the history
  • Loading branch information
ecederstrand committed Apr 25, 2018
1 parent 44a0b17 commit b7cec88
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 8 deletions.
18 changes: 10 additions & 8 deletions exchangelib/queryset.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,14 @@ def _additional_fields(self):
additional_fields.add(FieldPath(field=end_tz_field))
return additional_fields

def _format_items(self, items, return_format):
return {
self.VALUES: self._as_values,
self.VALUES_LIST: self._as_values_list,
self.FLAT: self._as_flat_values_list,
self.NONE: self._as_items,
}[return_format](items)

def _query(self):
from .folders import SHALLOW
from .items import Persona
Expand Down Expand Up @@ -300,13 +308,7 @@ def __iter__(self):

log.debug('Initializing cache')
_cache = []
result_formatter = {
self.VALUES: self._as_values,
self.VALUES_LIST: self._as_values_list,
self.FLAT: self._as_flat_values_list,
self.NONE: self._as_items,
}[self.return_format]
for val in result_formatter(self._query()):
for val in self._format_items(items=self._query(), return_format=self.return_format):
_cache.append(val)
yield val
self._cache = _cache
Expand Down Expand Up @@ -599,7 +601,7 @@ def iterator(self):
if self.is_cached:
return self._cache
# Return an iterator that doesn't bother with caching
return self._query()
return self._format_items(items=self._query(), return_format=self.return_format)

def get(self, *args, **kwargs):
""" Assume the query will return exactly one item. Return that item """
Expand Down
6 changes: 6 additions & 0 deletions tests/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3301,10 +3301,16 @@ def test_querysets(self):
[i.categories[0] for i in qs.only('categories').order_by('subject')],
[test_cat, test_cat, test_cat, test_cat]
)
# Test iterator
self.assertEqual(
set((i.subject, i.categories[0]) for i in qs.iterator()),
{('Item 0', test_cat), ('Item 1', test_cat), ('Item 2', test_cat), ('Item 3', test_cat)}
)
# Test that iterator() preserves the result format
self.assertEqual(
set((i[0], i[1][0]) for i in qs.values_list('subject', 'categories').iterator()),
{('Item 0', test_cat), ('Item 1', test_cat), ('Item 2', test_cat), ('Item 3', test_cat)}
)
self.assertEqual(qs.get(subject='Item 3').subject, 'Item 3')
with self.assertRaises(DoesNotExist):
qs.get(subject='Item XXX')
Expand Down

0 comments on commit b7cec88

Please sign in to comment.