Question

Is it possible in pymongo to have the collection.find() method return a custom cursor class that inherits from the base, but redefines how iteration occurs?

I’d like to instantiate app specific models from the mongo data within the cursor right at iteration time. The documents have a type attr that will determine what kind of instance should be created. I was thinking the next method could look within this data and decide which type to create and return. Inheriting from the cursor is easy enough, but I can’t figure out where to hook this into the find() operation?

EDIT Alternatively…

What I’m doing currently is to use yield to spit out a generator that will classify the object after the fact when the fetch is performed.

@classmethod
def gather(cls,place_id):
    """
    Gather instances of all the shouts for one place
    """
    shouts = cls.collection.find({'place_id':place_id})
    for s in shouts:
        yield Shout.classify(s)

@classmethod
def classify(cls,shout):
    if shout.type == 1:
        return Type1(dict(shout))
    elif shout.type == 2:
        return Type2(dict(shout))
    elif shout.type == 3:
        return Type3(dict(shout))

The problem is that this doesn’t preserve the original approach to bracket key access that is encapsulated in the default pymongo Cursor.

If I were to create a wrapper class that will just accept a cursor instance and wrap its methods, which magic methods will I need to override to preserve the iteration behavior of the original cursor? I’m thinking something like this:

class StuffCursor():

    cursor = False

    def __init__(self,cursor):
        self.cursor = cursor

    def __getattr__(self,attr):
        #This passes off most key-based calls, like bracket-access, to the cursor
        return getattr(self.cursor,attr)

This is just what I could think of, but anything that would stack some additional processing on top of an iterator, and then return the modified iterator, would work.

Answer 1

Hate to answer my own question, but this may prove to be of good use to anyone googling. Here’s what I ended up going with:

class ShoutCursor():
    """
    Custom cursor class that will instantiate shout models
    at query time.
    """

    cursor = False

    def __init__(self,cursor):
        self.cursor = cursor
        self.cursor.sort('created',DESCENDING)

    def __iter__(self):
        for ent in self.cursor:
            yield Shout.classify(ent)

    def __len__(self):
        return self.cursor.count()

    def __getitem__(self,index):
        try:
            ents = itertools.islice(self.cursor,index,index+1)
            return [Shout.classify(ent) for ent in ents][0]
        except TypeError:
            return list(Shout.classify(ent) for ent in itertools.islice(self.cursor,index.start,index.stop,index.step))

    def sort(*args,**kwargs):
        self.cursor.sort(*args,**kwargs)

    def __getattr__(self,attr):
        return getattr(self.cursor,attr)

The __iter__ and __getiem__ methods are where the modified, instantiated models are loaded up into generators are returned. The other methods preserve native cursor action. Passing a cursor into __init__ sets up the functionality, so special, per iteration logic can be performed as the objects are fetched from mongo.