Support generators as return values to __getitem__ on mongomock.collection.Cursor

[akellehe]

If a call to mongomock.collection.Cursor.__getitem__ expects a mongomock.collection.Cursor and gets a list (for example in mongoengine's slicing methods Model.objects.all()[start:finish]) we get an exception that list has no attribute/method next. The fix is super-easy, you just have to check here

https://github.com/vmalloc/mongomock/blob/master/mongomock/collection.py#L876

to see if the return value would be an instance of list:

    def __getitem__(self, index):
        arr = [x for x in self._dataset]
        count = len(arr)
        self._dataset = iter(arr)
        rval = arr[index]
        if isinstance(arr[index], list):
            return iter(arr[index])
        return arr[index]

There are a few other methods in the Cursor class that could use a little refactoring. distinct, for example, might be better as (for the same reason):

    def distinct(self, key):
        if not isinstance(key, basestring):
            raise TypeError('cursor.distinct key must be a string')
        unique = set()
        for x in iter(self._dataset):
            value = _resolve_key(key, x)
            if value == NOTHING: continue
            unique.add(value)
        self._dataset = (x for x in unique)
        return self

[Kimamisa]

Are you planning to create a PR ?

[abondis]

I have the issue too with lists. @akellehe fix to return an iterator works for me. I can do a PR, but I'm really not sure what other effects it would have. that said, running tox seemed to get success for all the tests.

[akellehe]

I'll submit a PR this evening.

[abondis]

I'm running with an issue with a complicated prepared queryset, with the return iter(arr[index]) fix

prepared_qs.first()
  File "mongoengine/mongoengine/queryset/base.py", line 262, in first
    queryset = self.clone()
  File "mongoengine/mongoengine/queryset/base.py", line 678, in clone
    return self.clone_into(self.__class__(self._document, self._collection_obj))
  File "mongoengine/mongoengine/queryset/base.py", line 700, in clone_into
    cls._cursor_obj = self._cursor_obj.clone()
AttributeError: 'listiterator' object has no attribute 'clone'
    cls._cursor_obj = self._cursor_obj.clone()
AttributeError: 'listiterator' object has no attribute 'clone'

I didn't identify why or how to reproduce simply, but returning a cursor (which has clone) fixes the problem in my case : return Cursor(self.collection, functools.partial(iter, arr[index]), self._limit)
Not sure this should be the solution as I still don't understand the whole internals of mongoengine and mongomock, but maybe that can help ?

[bcongdon]

Could @abondis' commit (2759170) be cherry picked into mongomock? I'm running into this issue, and it seems like the fix has been found, but not yet included.

(For those using mongoengine and running into issues using array-slicing on a queryset - i.e. results[:10] - I found that using #limit works, and has the same functionality - i.e. results.limit(10))

[pcorpet]

I think that @touilleMan did it in d7aeac1, thanks Emmanuel. please reopen if I'm wrong.