from datetime import datetime from logging import Logger import logging.config import asyncio import rollbot from rollbot.injection import ArgList, Lazy, Database, Data, Args, Arg, Subcommand logging.config.fileConfig('logging.conf', disable_existing_loggers=False) @rollbot.initialize_data class MyCounter: one: int = 0 two: int = 0 class MyBot(rollbot.Rollbot[str]): def read_config(self, key): return key def parse(self, raw): return rollbot.Message( origin_id="REPL", channel_id=".", sender_id=".", timestamp=datetime.now(), origin_admin=True, channel_admin=True, text=raw, attachments=[], ) async def respond(self, res): print(res, flush=True) async def goodbye_command(message, context): await context.respond(rollbot.Response.from_message(message, "Goodbye!")) async def count_command(message, context): args = message.text.split("count", maxsplit=1)[1].strip().split() name = args[0] if len(args) > 0 else "main" async with context.database() as db: await db.execute( "INSERT INTO counter VALUES (?, 1) \ ON CONFLICT (name) DO UPDATE SET count=count + 1", (name,), ) await db.commit() async with db.execute("SELECT count FROM counter WHERE name = ?", (name,)) as cursor: res = (await cursor.fetchone())[0] await context.respond(rollbot.Response.from_message(message, f"{name} = {res}")) @rollbot.as_command async def count2(name: Arg(0, required=False, default="main"), connect: Lazy(Database)): db = await connect() await db.execute( "INSERT INTO counter VALUES (?, 1) \ ON CONFLICT (name) DO UPDATE SET count=count + 2", (name,), ) await db.commit() async with db.execute("SELECT count FROM counter WHERE name = ?", (name,)) as cursor: res = (await cursor.fetchone())[0] return f"{name} = {res}" @rollbot.as_command async def count3(counter: Data(MyCounter).For(Args), store: Data(MyCounter), args: Args): counter.one += 1 counter.two += 2 await store.save(args, counter) return f"{args} = {counter}" @rollbot.as_command async def count6(counters: Data(MyCounter)): async for (key, counter) in counters.all(one=6): yield f"{key} = {counter}" async for (key, counter) in counters.all(two=6): yield f"{key} = {counter}" @rollbot.as_command async def lscounters(counters: Data(MyCounter)): async for (key, counter) in counters.all(): yield f"{key} = {counter}" @rollbot.as_command("count.") async def count_improved( subc: Subcommand, name: Subcommand.Arg(0, required=False, default="main"), counter: Data(MyCounter).For(Subcommand.Arg(0, required=False, default="main")), store: Data(MyCounter), log: Logger, ): if subc.name == "up": counter.one += 1 counter.two += 2 elif subc.name == "down": counter.one -= 1 counter.two -= 2 elif subc.name == "show": yield f"{name} = {counter}" log.info(f"Saving {counter} under {name}") await store.save(name, counter) @rollbot.on_startup async def make_table(context): async with context.database() as db: await db.execute( "CREATE TABLE IF NOT EXISTS counter ( \ name TEXT NOT NULL PRIMARY KEY, \ count INTEGER NOT NULL DEFAULT 0 \ );" ) await db.commit() @rollbot.on_shutdown async def shutdown(context): await context.respond(rollbot.Response(origin_id="REPL", channel_id=".", text="Shutting down!")) @rollbot.as_command def simple(): return "Simple!" @rollbot.as_command def generator(): yield "This is" yield "a generator!" @rollbot.as_command async def coroutine(): await asyncio.sleep(1.0) return "Here's a coroutine!" @rollbot.as_command async def asyncgen(): yield "This is" await asyncio.sleep(0.5) yield "an async" await asyncio.sleep(0.5) yield "generator!" config = rollbot.get_command_config().extend( rollbot.CommandConfiguration( commands={ "goodbye": goodbye_command, "count": count_command, }, call_and_response={ "hello": "Hello!", }, aliases={ "hi": "hello", "bye": "goodbye", }, bangs=("/",), ) ) bot = MyBot(config, "/tmp/my.db") async def run(): await bot.on_startup() try: while True: await bot.on_message(input("> ")) except EOFError: pass finally: await bot.on_shutdown() asyncio.run(run())