doccano/backend/examples/assignment/usecase.py

from typing import List

from django.shortcuts import get_object_or_404

from examples.assignment.strategies import StrategyName, create_assignment_strategy
from examples.models import Assignment, Example
from projects.models import Member, Project


def bulk_assign(project_id: int, strategy_name: StrategyName, member_ids: List[int], weights: List[int]) -> None:
    project = get_object_or_404(Project, pk=project_id)
    members = Member.objects.filter(project=project, pk__in=member_ids)
    if len(members) != len(member_ids):
        raise ValueError("Invalid member ids")
    # Sort members by member_ids
    members = sorted(members, key=lambda m: member_ids.index(m.id))

    unassigned_examples = Example.objects.filter(project=project, assignments__isnull=True)
    dataset_size = unassigned_examples.count()

    strategy = create_assignment_strategy(strategy_name, dataset_size, weights)
    assignments = strategy.assign()
    assignments = [
        Assignment(
            project=project,
            example=unassigned_examples[assignment.example],
            assignee=members[assignment.user].user,
        )
        for assignment in assignments
    ]
    Assignment.objects.bulk_create(assignments)