浏览代码

Implement first pass at simple decorators

Kirk Trombley 4 年之前
父节点
当前提交
765b78c3b7
共有 3 个文件被更改,包括 123 次插入14 次删除
  1. 1 0
      lib/rollbot/command/__init__.py
  2. 74 0
      lib/rollbot/command/decorators.py
  3. 48 14
      repl_driver.py

+ 1 - 0
lib/rollbot/command/__init__.py

@@ -1 +1,2 @@
+from .decorators import as_command, on_startup, on_shutdown, get_command_config
 from .failure import RollbotFailure

+ 74 - 0
lib/rollbot/command/decorators.py

@@ -0,0 +1,74 @@
+from collections.abc import Callable
+from typing import Union
+import inspect
+
+from ..types import Message, Context, CommandType, Response, StartupShutdownType, CommandConfiguration
+from .failure import RollbotFailureException
+
+decorated_startup: list[StartupShutdownType] = []
+decorated_shutdown: list[StartupShutdownType] = []
+decorated_commands: dict[str, CommandType] = {}
+
+
+def on_startup(fn):
+    decorated_startup.append(fn)
+    return fn
+
+
+def on_shutdown(fn):
+    decorated_shutdown.append(fn)
+    return fn
+
+
+def as_command(arg: Union[str, Callable]):
+    def impl(name, fn):
+        if inspect.isasyncgenfunction(fn):
+            lifted = fn
+        elif inspect.iscoroutinefunction(fn):
+            async def lifted(*args):
+                yield await fn(*args)
+        elif inspect.isgeneratorfunction(fn):
+            async def lifted(*args):
+                for res in fn(*args):
+                    yield res
+        elif inspect.isfunction(fn):
+            async def lifted(*args):
+                yield fn(*args)
+        else:
+            raise ValueError  # TODO details
+
+        async def command_impl(message: Message, context: Context):
+            args = []  # TODO implement dep injection
+
+            try:
+                async for result in lifted(*args):
+                    if isinstance(result, Response):
+                        response = result
+                    elif isinstance(result, str):
+                        response = Response.from_message(message, text=result)
+                    # TODO handle attachments, other special returns
+                    else:
+                        response = Response.from_message(message, str(result))
+                    await context.respond(response)
+            except RollbotFailureException as exc:
+                # TODO handle errors more specifically
+                await context.respond(Response.from_message(message, str(exc.failure)))
+
+        decorated_commands[name] = command_impl
+        return fn
+
+    if isinstance(arg, str):
+        return lambda fn: impl(arg, fn)
+    else:
+        return impl(arg.__name__, arg)
+
+
+def get_command_config() -> CommandConfiguration:
+    return CommandConfiguration(
+        commands=decorated_commands,
+        call_and_response={},
+        aliases={},
+        bangs=(),
+        startup=decorated_startup,
+        shutdown=decorated_shutdown,
+    )

+ 48 - 14
repl_driver.py

@@ -41,6 +41,7 @@ async def count_command(message, context):
     await context.respond(rollbot.Response.from_message(message, f"{name} = {res}"))
 
 
+@rollbot.on_startup
 async def make_table(context):
     async with context.database() as db:
         await db.execute("CREATE TABLE IF NOT EXISTS counter ( \
@@ -50,22 +51,55 @@ async def make_table(context):
         await db.commit()
 
 
-config = rollbot.CommandConfiguration(
-    commands={
-        "goodbye": goodbye_command,
-        "count": count_command,
-    },
-    call_and_response={
-        "hello": "Hello!",
-    },
-    aliases={
-        "hi": "hello",
-        "bye": "goodbye",
-    },
-    bangs=("/",),
-    startup=[make_table],
+@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():