tallybot/tallybot.py

#!/usr/bin/env python3
'''small bot that tallies up scores for strings that are incremented or
decremented in its channels and saves them to a sqlite database'''
import argparse
import logging
from pbnj.bot import Bot
from pbnj.logger import ColorFormatter
from collections import defaultdict
import sqlite3

INITIAL_CHANNELS = [
    '#meta',
    '#team',
    '#bots',
    '#_'
]

OTHER_BOTS = [
    'NormalHuman555',
    'zaphod',
    'cirno',
    'dustbotoBabilibot',
    'minerbot2',
    'sotdbot'
]

def setup_db(path):
    db_connection = sqlite3.connect(path, isolation_level=None)
    cur = db_connection.cursor()
    cur.execute('create table if not exists Votes(Message text not null, Up integer not null, Down integer not null)')
    cur.execute('create table if not exists Channels(Name text unique not null)')
    return (db_connection, cur)


def vote_up(message, cursor):
    cursor.execute("select Up from Votes where Message = ?", (message, ))
    result = cursor.fetchone()
    if result:
        cursor.execute('update Votes set Up = Up + 1 where Message = ?', (message, ))
    else:
        cursor.execute('insert into Votes values(?, ?, ?)', (message, 1, 0))


def vote_down(message, cursor):
    cursor.execute('select Down from Votes where Message = ?', (message, ))
    result = cursor.fetchone()
    if result:
        cursor.execute('update Votes set Down = Down + 1 where Message = ?', (message, ))
    else:
        cursor.execute('insert into Votes values(?, ?, ?)', (message, 0, 1))


def summarize(message, cursor):
    cursor.execute('select Message, Up, Down from Votes where Message = ?', (message, ))
    result = cursor.fetchone()
    count_votes = result[1] + result[2]
    total = result[1] - result[2]
    percent_positive = round(result[1]/count_votes*100, 2)
    return 'voted {}: {} votes total, {}% positive'.format(
        total,
        count_votes,
        percent_positive
    )


def recognize_inc_expr(message):
    valid_exprs = ['--', '++']
    if message.type == 'PRIVMSG':
        # Ignore 'https?://', tildegit.org/delucks/tallybot/issues/2
        if 'http://' in message.message or 'https://' in message.message:
            return False
        for expr in valid_exprs:
            # Only allow incr/decr for messages that have a space AFTER the operator, or messages that end in ++/--
            if expr+' ' in message.message or message.message.endswith(expr):
                return True
    return False

bot = Bot('karmabot', builtin_prefix='karmabot: ', use_builtin=True, connect_wait=5, initial_channels=INITIAL_CHANNELS, ignore={'nicks': OTHER_BOTS})
DBFILE = 'tallies.sqlite3'
db, cursor = setup_db(DBFILE)

@bot.command('karmabot: topfive')
def topfive(message):
    '''show current votes'''
    cursor.execute('select Message, Up-Down as Score from Votes order by Score desc limit(5)')
    for topic in cursor.fetchall():
        yield '{}: {}'.format(topic[0], topic[1])

@bot.command('karmabot: bottomfive')
def bottomfive(message):
    '''show current votes'''
    cursor.execute('select Message, Up-Down as Score from Votes order by Score asc limit(5)')
    for topic in cursor.fetchall():
        yield '{}: {}'.format(topic[0], topic[1])

#@bot.command('karmabot: help')
#def help(message):
#    '''Display this help message'''
#    return 'Usage: foo++/foo--, "karmabot: {top,bottom}five" for current scores, "karmabot: join #channel", and "karmabot: help"'

#@bot.command('karmabot: join')
#def join(message):
#    '''send me to other channel(s)! Comma-separate channel names'''
#    if len(message.args) < 1:
#        return 'Missing channel name! Try "join #channelname"'
#    else:
#        return bot.joinall(message.args)

@bot.command(recognize_inc_expr)
def increment(message):
    '''tallies up score for a certain string'''
    valid_exprs = {'--': False, '++': True}
    topic = None
    for expr in valid_exprs:
        if expr in message.message:
            topic = message.message.split(expr)[0].strip()
            if valid_exprs[expr]:
                vote_up(topic, cursor)
            else:
                vote_down(topic, cursor)
            break
    return '{}: {}'.format(topic, summarize(topic, cursor))

@bot.command('^!(botlist|rollcall)')
def botlist(message):
    '''Respond to required botlist command'''
    return 'Maintainer: delucks | Keeps a persistent tally of scores for strings mentioned in channels | Source: https://tildegit.org/delucks/tallybot | Use ++ or -- after a string'

if __name__ == '__main__':
    # Logging setup
    log = logging.getLogger('pbnj')
    sh = logging.StreamHandler()
    sh.setFormatter(ColorFormatter())
    log.addHandler(sh)
    # Argparsing
    p = argparse.ArgumentParser(description='tallybot')
    p.add_argument('--debug', action='store_true', help='enable debug logging')
    args = p.parse_args()
    log_lvl = logging.DEBUG if args.debug else logging.INFO
    log.setLevel(log_lvl)
    # Go go go
    bot.connect('127.0.0.1', 6667)
    bot.run()
    db.close()