rollbot4/repl_driver.py

from datetime import datetime
from logging import Logger
import asyncio

import rollbot
from rollbot.injection import ArgList, Lazy, Database, Data, Args, Arg, Subcommand


@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())