schneiderbot-tg/db.py

"""DB interface"""
import sqlite3

class Db():
    """ Abstract interface for db actions """
    def __init__(self):
        self.conn = sqlite3.connect("deaths.db")
        table_exist = self.conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='death'")
        if table_exist.fetchall() == []:
            self.conn.execute("CREATE TABLE death(name text, chat integer )")
            self.conn.execute("create index chat_idx on death(chat)")

    def __del__(self):
        self.conn.commit()

    def is_dead(self, user, chat):
        """Reports whether a user is already dead"""
        res = self.conn.execute("select * from death where chat = ? and name = ?", [chat, user])
        return res.fetchall() != []

    def kill(self, user, chat):
        """Kills a user"""
        self.conn.execute("insert into death(name, chat) values(?, ?)", [user, chat])

    def revive(self, user, chat):
        """Revives a user"""
        self.conn.execute("delete from death where chat = ? and name = ?", [chat, user])

    def get_dead_bodies(self, chat):
        """List the dead"""
        names = []
        for name in self.conn.execute("select * from death where chat = ?", [chat]).fetchall():
            names.append(name[0])
        return ", ".join(names)