hanabi-league/games_db_interface.py

from typing import List

import psycopg2.extras
from database import conn_manager

import hanabi.hanab_game


def store_actions(game_id: int, actions: List[hanabi.hanab_game.Action]):
    vals = []
    for turn, action in enumerate(actions):
        vals.append((game_id, turn, action.type.value, action.target, action.value or 0))

    conn = conn_manager.get_connection()
    cur = conn.cursor()
    psycopg2.extras.execute_values(
        cur,
        "INSERT INTO game_actions (game_id, turn, type, target, value) "
        "VALUES %s "
        "ON CONFLICT (game_id, turn) "
        "DO NOTHING",
        vals
    )
    conn.commit()


def store_deck_for_seed(seed: str, deck: List[hanabi.hanab_game.DeckCard]):
    vals = []
    for index, card in enumerate(deck):
        vals.append((seed, index, card.suitIndex, card.rank))

    conn = conn_manager.get_connection()
    cur = conn.cursor()
    psycopg2.extras.execute_values(
        cur,
        "INSERT INTO seeds (seed, card_index, suit_index, rank) "
        "VALUES %s "
        "ON CONFLICT (seed, card_index) "
        "DO NOTHING",
        vals
    )
    conn.commit()


def load_actions(game_id: int) -> List[hanabi.hanab_game.Action]:
    cur = conn_manager.get_new_cursor()
    cur.execute("SELECT type, target, value FROM game_actions "
                "WHERE game_id = %s "
                "ORDER BY turn ASC",
                (game_id,))
    actions = []
    for action_type, target, value in cur.fetchall():
        actions.append(
            hanabi.hanab_game.Action(hanabi.hanab_game.ActionType(action_type), target, value)
        )
    return actions


def load_deck(seed: str) -> List[hanabi.hanab_game.DeckCard]:
    cur = conn_manager.get_new_cursor()
    cur.execute("SELECT card_index, suit_index, rank FROM seeds "
                "WHERE seed = %s "
                "ORDER BY card_index ASC",
                (seed,)
                )
    deck = []
    for index, (card_index, suit_index, rank) in enumerate(cur.fetchall()):
        assert index == card_index
        deck.append(
            hanabi.hanab_game.DeckCard(suit_index, rank, card_index)
        )
    return deck