2026-09-07 19:43:48 +03:00

98 lines
3.3 KiB
Python

from __future__ import annotations
from dataclasses import dataclass
from django.db import transaction
from django.db.models import Count
from django.utils import timezone
from .models import Poll, PollOption, PollVote
from .utils import user_can_access_chat
@dataclass
class PollActionError(Exception):
message: str
status_code: int = 400
def __str__(self):
return self.message
def poll_is_effectively_closed(poll: Poll) -> bool:
return bool(poll.is_closed or (poll.closes_at and poll.closes_at <= timezone.now()))
def cast_poll_vote(poll: Poll, user, option_ids) -> tuple[Poll, list[int]]:
"""Replace a user's poll selection and refresh denormalized counters."""
try:
normalized_ids = list(dict.fromkeys(int(option_id) for option_id in option_ids))
except (TypeError, ValueError):
raise PollActionError("Некорректные варианты ответа") from None
if not normalized_ids:
raise PollActionError("Выберите хотя бы один вариант")
if not user_can_access_chat(poll.message.chat, user):
raise PollActionError("Нет доступа к этому голосованию", 403)
closed_by_deadline = False
options = []
with transaction.atomic():
locked_poll = (
Poll.objects.select_for_update()
.select_related("message", "message__chat")
.get(pk=poll.pk)
)
if poll_is_effectively_closed(locked_poll):
if not locked_poll.is_closed:
locked_poll.close()
closed_by_deadline = True
else:
if not locked_poll.is_multiple_choice and len(normalized_ids) > 1:
raise PollActionError("Only one option allowed")
options = list(
locked_poll.options.filter(pk__in=normalized_ids).order_by(
"position",
"id",
)
)
if len(options) != len(normalized_ids):
raise PollActionError("Некорректные варианты ответа")
PollVote.objects.filter(poll=locked_poll, voter=user).delete()
PollVote.objects.bulk_create(
[
PollVote(poll=locked_poll, option=option, voter=user)
for option in options
]
)
vote_counts = {
row["option_id"]: row["count"]
for row in (
PollVote.objects.filter(poll=locked_poll)
.values("option_id")
.annotate(count=Count("id"))
)
}
poll_options = list(locked_poll.options.all())
for option in poll_options:
option.vote_count = vote_counts.get(option.id, 0)
PollOption.objects.bulk_update(poll_options, ["vote_count"])
locked_poll.total_voters = (
PollVote.objects.filter(poll=locked_poll)
.values("voter_id")
.distinct()
.count()
)
locked_poll.save(update_fields=["total_voters", "updated_at"])
if closed_by_deadline:
raise PollActionError("Poll is closed")
return locked_poll, [option.id for option in options]