David Wolever
David Wolever

Reputation: 154494

Automatically annotate Django models retrieved through a ManyToMany field

I have a Student model that's related to a Course model through a ManyToMany field:

class Student(Model):
    name = TextField()

class Course(Model):
    name = TextField()
    students = ManyToManyField(Student, through="StudentCourse")

class StudentCourse(Model):
    student = ForeignKey(Student)
    course = ForeignKey(Course)
    section_name = TextField()

How can I automatically annotate students retrieved through the Course.students many-to-many field with their section in that course?

For example, instead of having to add the explicit extra on each query:

>>> students = course.students.extra(select={
...     "section_name": "app_student_course.section_name",
... })
>>> print students[0].section_name
u'First Section'

I can just:

>>> students = course.students.all()
>>> print students[0].section_name
u'First Section'

Thanks!

Upvotes: 6

Views: 1782

Answers (3)

ajaest
ajaest

Reputation: 627

Is it possible to replace the relation manager as follows. From there you can do whatever you want with the queryset:

from django.db.models.query import F

class Student(Model):
    name = TextField()

class Course(Model):
    name = TextField()
    students = ManyToManyField(Student, through="StudentCourse")

class StudentCourse(Model):
    # Set related_names so that it is easy to refer to the relation
    # with the through model
    student = ForeignKey(Student, related_name='student_courses')
    course = ForeignKey(Course, related_name='student_courses')
    section_name = TextField()

# Create a new customized manager
class StudentRelatedWithCourseManager(
    Course.students.related_manager_cls
):

    def get_queryset(self):
        # Gets the queryset of related Students ... 
        qs = super(StudentRelatedWithCourseManager, self)\
            .get_queryset()

        # Annotate it before is handled by the ManyRelatedDescriptor
        return qs.annotate(
            section_name=F('student_courses__section_name')
        )

# Replace the stock manager with your custom manager
Course.students.related_manager_cls = \
    StudentRelatedWithCourseManager

Upvotes: 5

Todor
Todor

Reputation: 16010

I would say the problem here is very similar with the one here, so my solution is almost identical to the one provided there.

I think prefetch_related with custom Prefetch is more suitable for such kind of problems instead of .annotations or .extra clauses. The benefit is that you get the whole related object instead of a single bit of it (so you can use more metadata), and there is zero performance drawback, only bigger memory footprint which can cause problems if its used for a large set of objects.

class StudenQuerySet(models.QuerySet):
    def prefetch_sections_for_course(self, course):
        return self.prefetch_related(models.Prefetch('studentcourse_set',
            queryset=StudentCourse.objects.filter(course=course),
            to_attr='course_sections'
        ))

class Student(Model):
    name = TextField()

    objects = StudenQuerySet.as_manager()

    @property
    def course_section_name(self):
        try:
            return self.course_sections[0].section_name
        except IndexError:
            #This student is not related with the prefetched course
            return None
        except AttributeError:
            raise AttributeError('You forgot to prefetch course sections')
            #or you can just
            return None

#usage
students = course.students.all().prefetch_sections_for_course(course)
for student in students:
    print student.course_section_name

Upvotes: 3

Spencer Judd
Spencer Judd

Reputation: 409

Why not just query the StudentCourse model? For example:

StudentCourse.objects.filter(course=course)

If you're iterating over StudentCourse objects and querying the student property, use select_related to make sure that you're not doing an additional DB Query for each Student.

student_courses = StudentCourse.objects.select_related('student').filter(course=course)
for sc in student_courses:
    print '%s: %s' % (sc.student.name, sc.section_name)

Upvotes: 0

Related Questions