WMICraft/logic/knights_queue.py

import random
from collections import deque


class KnightsQueue:
    def __init__(self, blue_team, red_team):
        self.queues = deque(blue_team), deque(red_team)
        self.team_idx_turn = random.randint(0, 1)

    def dequeue_knight(self):
        if self.both_teams_alive():
            knight = self.queues[self.team_idx_turn].popleft()
            if knight.max_hp <= 0:
                return self.dequeue_knight()
            else:
                self.queues[self.team_idx_turn].append(knight)
                self.next_turn()
                return knight
        raise Exception('Game has just ended')

    def both_teams_alive(self):
        return len(self.queues[0]) > 0 and len(self.queues[1]) > 0

    def next_turn(self):
        self.team_idx_turn = (self.team_idx_turn + 1) % 2