doccano/backend/data_export/pipeline/comments.py

import abc
from collections import defaultdict
from typing import Dict, List, Tuple

from django.db.models import QuerySet

from data_export.models import ExportedComment, ExportedExample


class Comments(abc.ABC):
    comment_class = ExportedComment
    column = "Comments"
    fields: Tuple[str, ...] = ("example", "user")  # To boost performance

    def __init__(self, examples: QuerySet[ExportedExample], user=None):
        self.comment_groups = defaultdict(list)
        comments = self.comment_class.objects.filter(example__in=examples)
        if user:
            comments = comments.filter(user=user)
        for comment in comments.select_related(*self.fields):
            self.comment_groups[comment.example.id].append(comment)

    def find_by(self, example_id: int) -> Dict[str, List[ExportedComment]]:
        return {self.column: self.comment_groups[example_id]}