Created
March 3, 2011 21:12
-
-
Save craigds/853579 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
=== modified file 'mptt/managers.py' | |
--- mptt/managers.py 2008-12-05 01:37:34 +0000 | |
+++ mptt/managers.py 2008-12-05 19:08:24 +0000 | |
@@ -11,15 +11,11 @@ | |
qn = connection.ops.quote_name | |
COUNT_SUBQUERY = """( | |
- SELECT COUNT(*) | |
- FROM %(rel_table)s | |
- WHERE %(mptt_fk)s = %(mptt_table)s.%(mptt_pk)s | |
+ %(mptt_fk)s = %(mptt_table)s.%(mptt_pk)s | |
)""" | |
CUMULATIVE_COUNT_SUBQUERY = """( | |
- SELECT COUNT(*) | |
- FROM %(rel_table)s | |
- WHERE %(mptt_fk)s IN | |
+ %(mptt_fk)s IN | |
( | |
SELECT m2.%(mptt_pk)s | |
FROM %(mptt_table)s m2 | |
@@ -29,6 +25,7 @@ | |
) | |
)""" | |
+ | |
class TreeManager(models.Manager): | |
""" | |
A manager for working with trees of objects. | |
@@ -47,7 +44,7 @@ | |
self.tree_id_attr = tree_id_attr | |
self.level_attr = level_attr | |
- def add_related_count(self, queryset, rel_model, rel_field, count_attr, | |
+ def add_related_count(self, queryset, rel_queryset, rel_field, count_attr, | |
cumulative=False): | |
""" | |
Adds a related item count to a given ``QuerySet`` using its | |
@@ -56,9 +53,8 @@ | |
Arguments: | |
- ``rel_model`` | |
- A ``Model`` class which has a relation to this `Manager``'s | |
- ``Model`` class. | |
+ ``rel_queryset`` | |
+ A ``QuerySet`` class for related model. | |
``rel_field`` | |
The name of the field in ``rel_model`` which holds the | |
@@ -74,24 +70,24 @@ | |
descendants, otherwise it will be for each item itself. | |
""" | |
opts = self.model._meta | |
+ rel_model = rel_queryset.model | |
if cumulative: | |
- subquery = CUMULATIVE_COUNT_SUBQUERY % { | |
- 'rel_table': qn(rel_model._meta.db_table), | |
+ rel_queryset = rel_queryset.extra(where=[CUMULATIVE_COUNT_SUBQUERY % { | |
'mptt_fk': qn(rel_model._meta.get_field(rel_field).column), | |
'mptt_table': qn(opts.db_table), | |
'mptt_pk': qn(opts.pk.column), | |
'tree_id': qn(opts.get_field(self.tree_id_attr).column), | |
'left': qn(opts.get_field(self.left_attr).column), | |
'right': qn(opts.get_field(self.right_attr).column), | |
- } | |
+ }]) | |
else: | |
- subquery = COUNT_SUBQUERY % { | |
- 'rel_table': qn(rel_model._meta.db_table), | |
+ rel_queryset = rel_queryset.extra(where=[COUNT_SUBQUERY % { | |
'mptt_fk': qn(rel_model._meta.get_field(rel_field).column), | |
'mptt_table': qn(opts.db_table), | |
'mptt_pk': qn(opts.pk.column), | |
- } | |
- return queryset.extra(select={count_attr: subquery}) | |
+ }]) | |
+ rel_queryset.query.add_count_column() | |
+ return queryset.extra(select={count_attr: rel_queryset.query}) | |
def get_query_set(self): | |
""" |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment