Counting all items within django-treebeard Materialised Path nodes

Question

I'm using django-treebeard to create a tree of Categories, each of which contains some Books:

from django.db import models
from treebeard.mp_tree import MP_Node

class Book(models.Model):
    title = models.CharField(max_length=255, blank=False, null=False)
    categories = models.ManyToManyField("Category", related_name="books", blank=True)

class Category(MP_Node):
    title = models.CharField(max_length=255, blank=False, null=False)
    book_count = models.IntegerField(default=0, blank=False,
        help_text="Number of Books in this Category")
    total_book_count = models.IntegerField(default=0, blank=False,
        help_text="Number of Books in this Category and its descendants")

I want to display the number of Books in each Category, stored in Category.book_count, so I have a post_save signal like this:

from django.dispatch import receiver

@receiver(post_save, sender=Book)
def book_post_save(sender, **kwargs):
    for category in kwargs["instance"].categories.all():
        category.book_count = category.books.count()
        category.save()

(The reality is a bit more complicated, but this is the gist.)

This works fine, and I can also do similar on an m2m_changed signal for if a Book's Categories have changed.

BUT, I also want to calculate total_book_count - the total number of Books in a Category and in all of its descendant Categories. I can do this with a method like this on Category:

def set_total_book_count(self):
    # All the Books in this Category's descendants:
    descendant_books = Book.objects.filter(categories__in=self.get_descendants())
    # Combine with the Books in this Category but avoid counting duplicates:
    books = (descendant_books | self.books).distinct()
    self.total_book_count = books.count()

This also works, and I can call it from signals. BUT this is all getting complicated, especially if I move a Category - I'd have to recalculate the total_book_count for the moved Category, all of its previous ancestors, and all of its new ancestors. At which point, maybe it's simpler to recalculate the counts for ALL Categories... a process that has its own difficulties (e.g. where to start recalculating).

Other people must have managed something similar and so I wonder if I'm badly re-inventing the wheel here? Are there existing implementations of this kind of thing?

score 0 · Answer 1 · answered Jun 15 '22 at 16:32

This is a partial answer, in that it's a simple way of recalculating all of the counts for all of the Categories. But I expect it could get slow if there are a lot of Categories and a lot of Books, given what set_total_book_count() has to do for every Category.

categories = Category.get_tree()

for category in categories:
    category.book_count = category.books.count()
    category.set_total_book_count()
    category.save()

Counting all items within django-treebeard Materialised Path nodes

1 Answers1