models.py

from django.db import models


class AnnotateByMixin():

    """Mixin that allows adding predefined annotations to queryset.

    Example:
        class DialogQuerySet(AnnotateByMixin, models.QuerySet):
            predefined_annotations = {
                'message_count': models.Count('message'),
            }
        class Dialog(models.Model):
            objects = DialogQuerySet.as_manager()
        annotated_dialogs = Dialog.objects.annotate_by('message_count')
        long_dialogs = annotated_dialogs.filter('message_count__gt=50')

    Attributes:
        predefined_annotations (dict): Predefined annotations.

    """

    predefined_annotations = {}

    def annotate_by(self, *args):
        """Annotate current queryset with predefined annotations.

        Args:
            *args: Annotation names.

        Returns:
            QuerySet: New queryset with annotations.

        """
        annotations = {field: self.predefined_annotations[field] for field in args}
        return self.annotate(**annotations)


class DialogQuerySet(AnnotateByMixin, models.QuerySet):
    predefined_annotations = {
        'message_count': models.Count('message'),
    }


class ChatUser(models.Model):
    name = models.CharField(max_length=63)
    companions = models.ManyToManyField('self', through='Dialog')


class Dialog(models.Model):
    client = models.ForeignKey('ChatUser', on_delete=models.CASCADE)
    servant = models.ForeignKey('ChatUser', on_delete=models.CASCADE)
    objects = DialogQuerySet.as_manager()


class Message(models.Model):
    dialog = models.ForeignKey('Dialog', on_delete=models.CASCADE)
    is_from_servant = models.BooleanField()
    text = models.CharField(max_length=127)