python - 如何在 Django 中向查询集添加自定义属性？

Question

考虑以下模型:

class User(AbstractUser):
    def messages(self):
        return Message.objects.filter(user=self).order_by('-created_on')

class Message(models.Model):
    user = models.ForeignKey('account.User', related_name='message', on_delete=models.CASCADE)
    created_on = models.DateTimeField(auto_now=True)
    data = models.TextField()
    read = models.BooleanField(default=False)

    objects = MessageManager()

我想重写 Message 模型的 QuerySet 行为以包含自定义方法 unread_count:

class MessageManager(models.Manager):
    def get_queryset(self):
        queryset = super().get_queryset()
        unread_count = queryset.filter(read=False).count()
        print(unread_count) # this works

        # these don't work
        queryset.unread_count = unread_count
        setattr(queryset, 'unread_count', unread_count)
        #

        return queryset

messages 方法工作正常，并且管理器内的 print 语句执行，但如果我尝试访问 unread_count，我会收到 AttributeError:

>>>user = User.objects.get(id=1)
>>>messages = user.messages()
2 # this is from the print statement
>>>messages.unread_count
AttributeError: 'QuerySet' object has no attribute 'unread_count'

我在这里缺少什么？

Answer 1

您应该对 models.QuerySet 进行子类化，而不是 models.Manager :

class MessageQuerySet(models.QuerySet):
    @property
    def unread_count(self):
        return self.filter(read=False).count()

然后使用as_manager模型中的方法:

class Message(models.Model):
    ...

    objects = MessageQuerySet.as_manager()

如果您需要向从管理器返回的 QuerySet 以及管理器本身添加方法，则需要使用 from_queryset :

class MessageManager(models.Manager):
    def some_manager_only_method(self):
        print('do stuff')

class MessageQuerySet(models.QuerySet):
    @property
    def unread_count(self):
        return self.filter(read=False).count()

class Message(models.Model):
    ...

    objects = MessageManager.from_queryset(MessageQuerySet)()