django-mptt/django-mptt — Issue #172: MPTT should honor custom QuerySets on subclassed TreeManager
Author: chrisdpratt
Created: 2012-02-23 16:36:41+00:00
State: open

If you subclass a `Manager` to add new methods, these methods are not chainable unless you also subclass `QuerySet`. Django uses the following pattern, which I've borrowed for my own projects:

class CustomQuerySet(models.query.QuerySet):
def some_new_method(self):
# do something here

class CustomManager(models.Manager):
def get_query_set(self):
return CustomQuerySet(self.model)

def some_new_method(self, *args, **kwargs):
return self.get_query_set().some_new_method(*args, **kwargs)

In this way, `some_new_method` is now chainable instead of only able to be called directly on the manager. However, if I do something similar, subclassing `TreeManager`, it doesn't work the majority of the time because all of the methods on the model, such as `get_ancestors`, `get_descendants()`, `get_siblings`, etc. do not use `_default_manager` but `_tree_manager`, which is added to the model via the following lines in mptt/models.py:

# make sure we have a tree manager somewhere
tree_manager = TreeManager()
tree_manager.contribute_to_class(cls, '_tree_manager')
tree_manager.init_from_model(cls)

# avoid using ManagerDescriptor, so instances can refer to self._tree_manager
setattr(cls, '_tree_manager', tree_manager)

The end result is that `TreeManager`, not the subclass actually being used on the model is always used, and returns a plain-jane `QuerySet`. I can see the logic behind doing this, i.e. ensuring that the code is working with an actual `TreeManager`, but sure there's a way that this constraint can be upheld without ignoring customizations. For all intents and purposes the `_default_manager` is completely ignored for most cases, even if it is itself a TreeManager. Instead, there should be some check to see if `_default_manager` is a `TreeManager` or a subclass of `TreeManager`, and only add a `TreeManager` if it's not present in some form. Off the top of my head, something like this (may not actually be optimal):

If cls._default_manager.__class__ is TreeManager or TreeManager in cls._default_manager.__class__.__bases__:
tree_manager = cls._default_manager
else:
# add the TreeManager as currently done

There's probably a better way than that, but the idea is to set `_tree_manager` to `_default_manager` if it's a type of `TreeManager`. That way, any customizations done in a subclass of `TreeManager` come along for the ride.