pymongo - 使用自定义游标类

Question

是否可以在 pymongo 中让collection.find() 方法返回一个从基类继承的自定义游标类，但重新定义了迭代的发生方式？

我想在迭代时从光标内的 mongo 数据实例化特定于应用程序的模型。这些文档有一个type attr，它将确定应该创建什么样的实例。我在想next 方法可以查看这些数据并决定创建和返回哪种类型。从光标继承很容易，但我不知道在哪里将它挂钩到 find() 操作？

编辑或者...

我目前正在做的是使用yield 吐出一个生成器，该生成器将在执行提取后对对象进行分类。

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

问题在于，这并没有保留封装在默认 pymongo Cursor 中的括号键访问的原始方法。

如果我要创建一个只接受一个游标实例并包装其方法的包装类，我需要重写哪些魔术方法以保留原始游标的迭代行为？我在想这样的事情：

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)

这只是我能想到的，但是任何会在迭代器之上堆叠一些额外处理，然后返回修改后的迭代器的方法都可以。

Answer 1

讨厌回答我自己的问题，但这可能对任何使用谷歌搜索的人都有好处。这是我最终的结果：

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)

__iter__ 和 __getiem__ 方法是将修改后的实例化模型加载到生成器中的位置返回。其他方法保留本机光标操作。将光标传递到__init__ 设置功能，如此特别，每次迭代逻辑都可以在从 mongo 获取对象时执行。

Answer 2

我不知道是否可以使用“猴子补丁”来实现这一点，但我最终做到了，如下所示。这允许我在 Cursor 类上添加/覆盖函数。

from pymongo import MongoClient
from pymongo.collection import Collection
from pymongo.cursor import Cursor
from pymongo.database import Database

class CustomCursor(Cursor):
    # Customize Cursor here, either by overriding or extending with more methods


class CustomCollection(Collection):
    def __init__(self, database: Database, name: str) -> None:
        super().__init__(database, name)

    # find_one seems to use find, but you may need to override in more places.
    def find(self, *args, **kwargs):
        return CustomCursor(self, *args, **kwargs)


class CustomDatabase(Database):
    def __init__(self, client: MongoClient, name: str) -> None:
        super().__init__(client, name)

    def __getitem__(self, name: str) -> CustomCollection:
        return CustomCollection(self, name)

# usage:
client = MongoClient('127.0.0.1', 27017)
database = CustomDatabase(client, 'mydb')

我用from_dict 和to_dict 方法定义了一个Serializable 基类。这让我可以在CustomCursor 上定义to_list()，这样我就可以查询并取回类实例：

result: List[School] = database.schools.find().to_list(School)