def _handle_unwind_stage(in_collection, unused_database, options):
if not isinstance(options, dict):
options = {'path': options}
path = options['path']
if not isinstance(path, six.string_types) or path[0] != '$':
raise ValueError(
'$unwind failed: exception: field path references must be prefixed '
"with a '$' '%s'" % path)
path = path[1:]
should_preserve_null_and_empty = options.get('preserveNullAndEmptyArrays')
include_array_index = options.get('includeArrayIndex')
unwound_collection = []
for doc in in_collection:
try:
array_value = helpers.get_value_by_dot(doc, path)
except KeyError:
if should_preserve_null_and_empty:
unwound_collection.append(doc)
continue
if array_value is None:
if should_preserve_null_and_empty:
unwound_collection.append(doc)
continue
if array_value == []:
if should_preserve_null_and_empty:
new_doc = copy.deepcopy(doc)
# We just ran a get_value_by_dot so we know the value exists.
helpers.delete_value_by_dot(new_doc, path)
unwound_collection.append(new_doc)
continue
if isinstance(array_value, list):
iter_array = enumerate(array_value)
else:
iter_array = [(None, array_value)]
for index, field_item in iter_array:
new_doc = copy.deepcopy(doc)
new_doc = helpers.set_value_by_dot(new_doc, path, field_item)
if include_array_index:
new_doc = helpers.set_value_by_dot(new_doc,
include_array_index, index)
unwound_collection.append(new_doc)
return unwound_collection
def test__set_value_by_dot(self):
"""Test set_value_by_dot"""
for doc, key, expected in (
({}, 'a', {'a': 42}),
({'a': 1}, 'a', {'a': 42}),
({'a': {'b': 1}}, 'a', {'a': 42}),
({'a': {'b': 1}}, 'a.b', {'a': {'b': 42}}),
({'a': [{'b': 1}]}, 'a.0', {'a': [42]}),
({'a': [{'b': 1}]}, 'a.0.b', {'a': [{'b': 42}]})):
ret = set_value_by_dot(doc, key, 42)
assert ret is doc
self.assertEqual(ret, expected)
