From 11b4ade690f7390e7091acd6da3ec4d1a2c0c8f1 Mon Sep 17 00:00:00 2001 From: Joscha Date: Fri, 27 Jul 2018 12:21:18 +0000 Subject: [PATCH] Port PlusOne to new yaboli version --- .gitignore | 5 +- database.py | 86 ++++++++++++++++++++++++++++++ plusone.py | 148 ++++++++++++++++++++++++++-------------------------- 3 files changed, 165 insertions(+), 74 deletions(-) create mode 100644 database.py diff --git a/.gitignore b/.gitignore index 5fb90f6..8fd98c7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ +**/__pycache__ yaboli -__pycache__ +websockets +join_rooms.py +*.cookie *.db diff --git a/database.py b/database.py new file mode 100644 index 0000000..7f9eb37 --- /dev/null +++ b/database.py @@ -0,0 +1,86 @@ +import asyncio +from functools import wraps +import sqlite3 +import threading + +__all__ = ["Database"] + + +def shielded(afunc): + #@wraps(afunc) + async def wrapper(*args, **kwargs): + return await asyncio.shield(afunc(*args, **kwargs)) + return wrapper + +class PooledConnection: + def __init__(self, pool): + self._pool = pool + + self.connection = None + + async def open(self): + self.connection = await self._pool._request() + + async def close(self): + conn = self.connection + self.connection = None + await self._pool._return(conn) + + async def __aenter__(self): + await self.open() + return self + + async def __aexit__(self, exc_type, exc, tb): + await self.close() + +class Pool: + def __init__(self, filename, size=10): + self.filename = filename + self.size = size + + self._available_connections = asyncio.Queue() + + for i in range(size): + conn = sqlite3.connect(self.filename, check_same_thread=False) + self._available_connections.put_nowait(conn) + + def connection(self): + return PooledConnection(self) + + async def _request(self): + return await self._available_connections.get() + + async def _return(self, conn): + await self._available_connections.put(conn) + +class Database: + def __init__(self, filename, pool_size=10, event_loop=None): + self._filename = filename + self._pool = Pool(filename, size=pool_size) + self._loop = event_loop or asyncio.get_event_loop() + + def operation(func): + @wraps(func) + @shielded + async def wrapper(self, *args, **kwargs): + async with self._pool.connection() as conn: + return await self._run_in_thread(func, conn.connection, *args, **kwargs) + return wrapper + + @staticmethod + def _target_function(loop, future, func, *args, **kwargs): + result = None + try: + result = func(*args, **kwargs) + finally: + loop.call_soon_threadsafe(future.set_result, result) + + async def _run_in_thread(self, func, *args, **kwargs): + finished = asyncio.Future() + target_args = (self._loop, finished, func, *args) + + thread = threading.Thread(target=self._target_function, args=target_args, kwargs=kwargs) + thread.start() + + await finished + return finished.result() diff --git a/plusone.py b/plusone.py index bd00477..6c7cd99 100644 --- a/plusone.py +++ b/plusone.py @@ -1,14 +1,16 @@ +import asyncio +import re + import yaboli from yaboli.utils import * -import asyncio -import random -import re -import sys +import database + +# List of rooms kept in separate file, which is .gitignore'd +import join_rooms - -class PointDB(yaboli.Database): - @yaboli.Database.operation +class PointDB(database.Database): + @database.Database.operation def initialize(conn): cur = conn.cursor() cur.execute(( @@ -18,21 +20,21 @@ class PointDB(yaboli.Database): ")" )) conn.commit() - - @yaboli.Database.operation + + @database.Database.operation def add_point(conn, nick): nick = mention_reduced(nick) cur = conn.cursor() - + cur.execute("INSERT OR IGNORE INTO Points (nick, points) VALUES (?, 0)", (nick,)) cur.execute("UPDATE Points SET points=points+1 WHERE nick=?", (nick,)) conn.commit() - - @yaboli.Database.operation + + @database.Database.operation def points_of(conn, nick): nick = mention_reduced(nick) cur = conn.cursor() - + cur.execute("SELECT points FROM Points WHERE nick=?", (nick,)) res = cur.fetchone() if res is not None: @@ -40,65 +42,51 @@ class PointDB(yaboli.Database): else: return 0 + +PLUSONE_RE = r"(\+1|:\+1:|:bronze(!\?|\?!)?:)\s*(.*)" +MENTION_RE = r"((for|to)\s+)?@(\S+)" + class PlusOne(yaboli.Bot): """ Count +1s awarded to users by other users. """ - - PLUSONE_RE = r"(\+1|:\+1:|:bronze(!\?|\?!)?:)\s*(.*)" - MENTION_RE = r"((for|to)\s+)?@(\S+)" - - def __init__(self, db): - super().__init__("PlusOne") - - self.db = db - - self.help_general = "/me counts :+1:s." - self.help_specific = ( + + async def created(self, room): + room.pointsdb = PointDB(f"points-{room.roomname}.db") + await room.pointsdb.initialize() + + async def send(self, room, message): + ping_text = ":bronze!?:" + short_help = "/me counts :+1:s" + long_help = ( "Counts +1/:+1:/:bronze:s: Simply reply \"+1\" to someone's message to award them a point.\n" "Alternatively, specify a person with: \"+1 [to|for] @person\"\n" + "\n" "!points - show your own points\n" - "!points [ ...] - list other people's points\n\n" - "Created by @Garmy using yaboli.\n" - "For additional info, try \"!help @{nick} \". Topics:\n" + "!points [ ...] - list other people's points\n" + "\n" + "Created by @Garmy using https://github.com/Garmelon/yaboli.\n" ) - self.help_specific += self.list_help_topics() - self.ping_message = ":bronze!?:" - - self.register_command("points", self.command_points, specific=False) - self.register_trigger(self.PLUSONE_RE, self.trigger_plusone) - - async def trigger_plusone(self, message, match): - nick = None - specific = re.match(self.MENTION_RE, match.group(3)) + await self.botrulez_ping_general(room, message, ping_text=ping_text) + await self.botrulez_ping_specific(room, message, ping_text=ping_text) + await self.botrulez_help_general(room, message, help_text=short_help) + await self.botrulez_help_specific(room, message, help_text=long_help) + await self.botrulez_uptime(room, message) + await self.botrulez_kill(room, message) + await self.botrulez_restart(room, message) - if specific: - nick = specific.group(3) - elif message.parent: - parent_message = await self.room.get_message(message.parent) - nick = parent_message.sender.nick - - if nick is None: - await self.room.send("You can't +1 nothing...", message.mid) - elif similar(nick, message.sender.nick): - newtext = random.choice([ - "Don't +1 yourself, that's... nasty.", - "Don't +1 yourself, that's... It just doesn't work that way, alright?", - "No points for " + nick + "!", - "No", # Attempt to trigger @Elements - "Would you kindly stop that?", - "Where'd you get that idea from? It would never work." - ]) - await self.room.send(newtext, message.mid) - else: - await self.db.add_point(nick) - await self.room.send(f"Point for user {mention(nick)} registered.", message.mid) - - async def command_points(self, message, argstr): + await self.command_points(room, message) + + await self.trigger_plusone(room, message) + + forward = send + + @yaboli.command("points", specific=False, args=True) + async def command_points(self, room, message, argstr): args = self.parse_args(argstr) if not args: - points = await self.db.points_of(message.sender.nick) - await self.room.send( + points = await room.pointsdb.points_of(message.sender.nick) + await room.send( f"You have {points} point{'s' if points != 1 else ''}.", message.mid ) @@ -107,21 +95,35 @@ class PlusOne(yaboli.Bot): for arg in args: if arg[:1] == "@": nick = arg[1:] - points = await self.db.points_of(nick) - response.append(f"@{mention(nick)} has {points} point{'' if points == 1 else 's'}.") else: - response.append(f"{arg!r} is not a mention.") - await self.room.send("\n".join(response), message.mid) + nick = arg + points = await room.pointsdb.points_of(nick) + response.append(f"{mention(nick)} has {points} point{'' if points == 1 else 's'}.") + await room.send("\n".join(response), message.mid) + + @yaboli.trigger(PLUSONE_RE) + async def trigger_plusone(self, room, message, match): + nick = None + specific = re.match(MENTION_RE, match.group(3)) + + if specific: + nick = specific.group(3) + elif message.parent: + parent_message = await room.get_message(message.parent) + nick = parent_message.sender.nick + + if nick is None: + await room.send("You can't +1 nothing...", message.mid) + elif similar(nick, message.sender.nick): + await room.send("Don't +1 yourself, that's not how things work.", message.mid) + else: + await room.pointsdb.add_point(nick) + await room.send(f"Point for user {mention(nick)} registered.", message.mid) def main(): - if len(sys.argv) == 3: - db = PointDB(sys.argv[2]) - asyncio.get_event_loop().run_until_complete(db.initialize()) - run_bot(PlusOne, sys.argv[1], db) - else: - print("USAGE:") - print(f" {sys.argv[0]} ") - return + bot = PlusOne("PlusOne") + join_rooms.join_rooms(bot) + asyncio.get_event_loop().run_forever() if __name__ == "__main__": main()