Ticket #17270: manager.5.py

from types import MethodType
from django.db import models
from django.db.models.query import QuerySet
from django.db.models.query import EmptyQuerySet

class ManagerQuerySetDecorator(models.Manager):
    """
    Extends the operation methods of the manager on subqueries QuerySet objects.
    
    How to use:
    
    class ManagerQuerySet(models.query.QuerySet):
        def get_fun_people(self):
            return self.filter(fun=True)
        
        def get_not_fun_people(self):
            return self.filter(fun=False)
        
    class PersonManager(models.Manager):
        def __init__(self):
            super(PersonManager, self).__init__()
            self.implement_queryset_cls(ManagerQuerySet)
            self.implement_queryset_methods('get_fun_people', 'get_not_fun_people')
            
    class Person(models.Model):
        first_name = models.CharField(max_length=30)
        last_name = models.CharField(max_length=30)
        fun = models.BooleanField()
        objects = PersonManager()
        objects_second = PersonManagerSecond(ManagerQuerySet, ['get_fun_people', 'get_not_fun_people'])
       
        def __unicode__(self):
            return u"%s %s" % (self.first_name, self.last_name)
    
    # This works without a decorator
    Person.objects.get_fun_people().filter(firstname='John')
    
    # The decorator also provides such a possibility, which does not work currently in django.
    Person.objects.filter(firstname='John').get_fun_people()
    """
    def __init__(self, queryset_cls=None, queryset_method_names=[]):
        super(ManagerQuerySetDecorator, self).__init__()
        self._queryset_cls = queryset_cls or QuerySet
        self.implement_queryset_methods(*queryset_method_names)

    def implement_queryset_cls(self, cls):
        self._queryset_cls = cls

    def implement_queryset_methods(self, *method_names):
        for method_name in method_names:
            self.implement_queryset_method(method_name)

    def implement_queryset_method(self, name):
        def queryset_method_deco(*args, **kwarg):
            return getattr(self.get_query_set(), name)(*args, **kwarg)
        setattr(self, name, queryset_method_deco)

    def get_empty_query_set(self):
        return self._queryset_cls(self.model, using=self._db).none()

    def get_query_set(self):
        return self._queryset_cls(self.model, using=self._db)