Source

pmxbot / pmxbot / karma.py

The default branch has multiple heads

Full commit
# vim:ts=4:sw=4:noexpandtab

from __future__ import print_function, absolute_import

import itertools
import re
import random

import pmxbot
from . import storage
from .core import command

class SameName(ValueError): pass
class AlreadyLinked(ValueError): pass


class Karma(storage.SelectableStorage):
	@classmethod
	def initialize(cls):
		cls.store = cls.from_URI(pmxbot.config.database)
		cls._finalizers.append(cls.finalize)

	@classmethod
	def finalize(cls):
		del cls.store

class SQLiteKarma(Karma, storage.SQLiteStorage):
	def init_tables(self):
		CREATE_KARMA_VALUES_TABLE = '''
			CREATE TABLE IF NOT EXISTS karma_values (karmaid INTEGER NOT NULL, karmavalue INTEGER, primary key (karmaid))
		'''
		CREATE_KARMA_KEYS_TABLE = '''
			CREATE TABLE IF NOT EXISTS karma_keys (karmakey varchar, karmaid INTEGER, primary key (karmakey))
		'''
		CREATE_KARMA_LOG_TABLE = '''
			CREATE TABLE IF NOT EXISTS karma_log (karmakey varchar, logid INTEGER, change INTEGER)
		'''
		self.db.execute(CREATE_KARMA_VALUES_TABLE)
		self.db.execute(CREATE_KARMA_KEYS_TABLE)
		self.db.execute(CREATE_KARMA_LOG_TABLE)
		self.db.commit()

	def lookup(self, thing):
		thing = thing.strip().lower()
		LOOKUP_SQL = 'SELECT karmavalue from karma_keys k join karma_values v on k.karmaid = v.karmaid where k.karmakey = ?'
		try:
			karma = self.db.execute(LOOKUP_SQL, [thing]).fetchone()[0]
		except:
			karma = 0
		if karma == None:
			karma = 0
		return karma

	def set(self, thing, value):
		thing = thing.strip().lower()
		value = int(value)
		UPDATE_SQL = 'UPDATE karma_values SET karmavalue = ? where karmaid = (select karmaid from karma_keys where karmakey = ?)'
		res = self.db.execute(UPDATE_SQL, (value, thing))
		if res.rowcount == 0:
			INSERT_VALUE_SQL = 'INSERT INTO karma_values (karmavalue) VALUES (?)'
			INSERT_KEY_SQL = 'INSERT INTO karma_keys (karmakey, karmaid) VALUES (?, ?)'
			ins = self.db.execute(INSERT_VALUE_SQL, [value])
			self.db.execute(INSERT_KEY_SQL, (thing, ins.lastrowid))
		self.db.commit()

	def change(self, thing, change):
		thing = thing.strip().lower()
		value = int(self.lookup(thing)) + int(change)
		UPDATE_SQL = 'UPDATE karma_values SET karmavalue = ? where karmaid = (select karmaid from karma_keys where karmakey = ?)'
		res = self.db.execute(UPDATE_SQL, (value, thing))
		if res.rowcount == 0:
			INSERT_VALUE_SQL = 'INSERT INTO karma_values (karmavalue) VALUES (?)'
			INSERT_KEY_SQL = 'INSERT INTO karma_keys (karmakey, karmaid) VALUES (?, ?)'
			ins = self.db.execute(INSERT_VALUE_SQL, [value])
			self.db.execute(INSERT_KEY_SQL, (thing, ins.lastrowid))
		self.db.commit()

	def list(self, select=0):
		KARMIC_VALUES_SQL = 'SELECT karmaid, karmavalue from karma_values order by karmavalue desc'
		KARMA_KEYS_SQL= 'SELECT karmakey from karma_keys where karmaid = ?'

		karmalist = self.db.execute(KARMIC_VALUES_SQL).fetchall()
		karmalist.sort(key=lambda x: int(x[1]), reverse=True)
		if select > 0:
			selected = karmalist[:select]
		elif select < 0:
			selected = karmalist[select:]
		else:
			selected = karmalist
		keysandkarma = []
		for karmaid, value in selected:
			keys = [x[0] for x in self.db.execute(KARMA_KEYS_SQL, [karmaid])]
			keysandkarma.append((keys, value))
		return keysandkarma

	def link(self, thing1, thing2):
		if thing1 == thing2:
			raise SameName("Attempted to link two of the same name")
		GET_KARMAID_SQL = 'SELECT karmaid FROM karma_keys WHERE karmakey = ?'
		try:
			t1id = self.db.execute(GET_KARMAID_SQL, [thing1]).fetchone()[0]
		except TypeError:
			raise KeyError(thing1)
		t1value = self.lookup(thing1)
		try:
			t2id = self.db.execute(GET_KARMAID_SQL, [thing2]).fetchone()[0]
		except TypeError:
			raise KeyError(thing2)
		if t1id == t2id:
			raise AlreadyLinked("Those two are already linked")
		t2value = self.lookup(thing2)

		newvalue = t1value + t2value
		# update the keys so t2 points to t1s value
		self.db.execute('UPDATE karma_keys SET karmaid = ? where karmaid = ?',
			(t1id, t2id))
		# drop the old value row for neatness
		self.db.execute('DELETE FROM karma_values WHERE karmaid = ?', (t2id,))
		# set the new combined value
		self.db.execute('UPDATE karma_values SET karmavalue = ? where karmaid = ?',
			(newvalue, t1id))
		self.db.commit()

	def _get(self, id):
		"""
		Return keys and value for karma id
		"""
		VALUE_SQL = "SELECT karmavalue from karma_values where karmaid = ?"
		KEYS_SQL = "SELECT karmakey from karma_keys where karmaid = ?"
		value = self.db.execute(VALUE_SQL, [id]).fetchall()[0][0]
		keys_cur = self.db.execute(KEYS_SQL, [id]).fetchall()
		keys = sorted(x[0] for x in keys_cur)
		return keys, value

	def search(self, term):
		query = "SELECT distinct karmaid from karma_keys where karmakey like ?"
		matches = (id for (id,) in self.db.execute(query, '%%'+term+'%%'))
		return (self._lookup(id) for id in matches)

	def export_all(self):
		return self.list()


class MongoDBKarma(Karma, storage.MongoDBStorage):
	collection_name = 'karma'
	def lookup(self, thing):
		thing = thing.strip().lower()
		res = self.db.find_one({'names':thing})
		return res['value'] if res else 0

	def set(self, thing, value):
		thing = thing.strip().lower()
		value = int(value)
		query = {'names': {'$in': [thing]}}
		oper = {'$set': {'value': value}, '$addToSet': {'names': thing}}
		self.db.update(query, oper, upsert=True)

	def change(self, thing, change):
		thing = thing.strip().lower()
		change = int(change)
		query = {'names': {'$in': [thing]}}
		oper = {'$inc': {'value': change}, '$addToSet': {'names': thing}}
		self.db.update(query, oper, upsert=True)

	def list(self, select=0):
		res = list(self.db.find().sort('value', storage.pymongo.DESCENDING))

		if select > 0:
			selected = res[:select]
		elif select < 0:
			selected = res[select:]
		else:
			selected = res
		aslist = lambda val: val if isinstance(val, list) else [val]
		return [
			(aslist(rec['names']), rec['value'])
			for rec in selected
		]

	def link(self, thing1, thing2):
		thing1 = thing1.strip().lower()
		thing2 = thing2.strip().lower()
		if thing1 == thing2:
			raise SameName("Attempted to link two of the same name")
		rec = self.db.find_one({'names': thing2})
		if thing1 in rec['names']:
			raise AlreadyLinked("Those two are already linked")
		if not rec: raise KeyError(thing2)
		try:
			query = {'names': thing1}
			update = {
				'$inc': {'value': rec['value']},
				'$pushAll': {'names': rec['names']},
			}
			self.db.update(query, update, safe=True)
		except Exception:
			raise KeyError(thing1)
		self.db.remove(rec)

	def search(self, term):
		pattern = re.compile('.*' + re.escape(term) + '.*')
		return (
			(rec['names'], rec['value'])
			for rec in self.db.find({'names': pattern})
		)

	def import_(self, item):
		names, value = item
		self.db.insert(dict(
			names = names,
			value = value,
			))

	def _all_names(self):
		return set(itertools.chain.from_iterable(
			names
			for names, value in self.search('')
		))

	def repair_duplicate_names(self):
		"""
		Prior to 1101.1.1, pmxbot would incorrectly create new karma records
		for individuals with multiple names.
		This routine corrects those records.
		"""
		for name in self._all_names():
			cur = self.db.find({'names': name})
			main_doc = next(cur)
			for duplicate in cur:
				query = {'_id': main_doc['_id']}
				update = {
					'$inc': {'value': duplicate['value']},
					'$pushAll': {'names': duplicate['names']},
				}
				self.db.update(query, update, safe=True)
				self.db.remove(duplicate)

@command("karma", aliases=("k",), doc="Return or change the karma value for "
	"some(one|thing)")
def karma(client, event, channel, nick, rest):
	karmee = rest.strip('++').strip('--').strip('~~')
	if '++' in rest:
		Karma.store.change(karmee, 1)
	elif '--' in rest:
		Karma.store.change(karmee, -1)
	elif '~~' in rest:
		change = random.choice([-1, 0, 1])
		Karma.store.change(karmee, change)
		if change == 1:
			return "%s karma++" % karmee
		elif change == 0:
			return "%s karma shall remain the same" % karmee
		elif change == -1:
			return "%s karma--" % karmee
	elif '==' in rest:
		t1, t2 = rest.split('==')
		try:
			Karma.store.link(t1, t2)
		except SameName:
			Karma.store.change(nick, -1)
			return "Don't try to link a name to itself!"
		except AlreadyLinked:
			return "Those names were previously linked."
		score = Karma.store.lookup(t1)
		return "%s and %s are now linked and have a score of %s" % (t1, t2, score)
	else:
		karmee = rest or nick
		score = Karma.store.lookup(karmee)
		return "%s has %s karmas" % (karmee, score)

@command("top10", aliases=("top",), doc="Return the top n (default 10) "
	"highest entities by Karmic value. Use negative numbers for the bottom "
	"N.")
def top10(client, event, channel, nick, rest):
	if rest:
		topn = int(rest)
	else:
		topn = 10
	selection = Karma.store.list(topn)
	res = ' '.join('(%s: %s)' % (', '.join(n), k) for n, k in selection)
	return res

@command("bottom10", aliases=("bottom",), doc="Return the bottom n (default "
	"10) lowest entities by Karmic value. Use negative numbers for the "
	"bottom N.")
def bottom10(client, event, channel, nick, rest):
	if rest:
		topn = -int(rest)
	else:
		topn = -10
	selection = Karma.store.list(topn)
	res = ' '.join('(%s: %s)' % (', '.join(n), k) for n, k in selection)
	return res