Add more util injectors, reformat, reorganize imports

lib/rollbot/__init__.py

@@ -2,4 +2,3 @@ from .bot import Rollbot
 from .failure import RollbotFailure
 from .types import *
 from .decorators import *
-from .injection import *

lib/rollbot/decorators/as_command.py

@@ -6,10 +6,11 @@ import inspect
 from ..types import (
     Message,
     Context,
+    Command,
     CommandType,
     Response,
 )
-from ..injection import Injector, inject_all, MessageInjector, ContextInjector
+from ..injection import Injector, inject_all, MessageInjector, ContextInjector, CommandInjector
 from .error_handling import with_failure_handling
 
 decorated_commands: dict[str, CommandType] = {}
@@ -51,6 +52,8 @@ def _get_injectors(fn: Callable[..., Any]) -> list[Injector]:
             injectors.append(MessageInjector)
         elif annot == Context:
             injectors.append(ContextInjector)
+        elif annot == Command:
+            injectors.append(CommandInjector)
         elif isinstance(annot, Injector):
             injectors.append(annot)
         else:
@@ -61,11 +64,10 @@ def _get_injectors(fn: Callable[..., Any]) -> list[Injector]:
 def _make_response(message: Message, result: Any) -> Response:
     if result is None or isinstance(result, Response):
         return result
-    elif isinstance(result, str):
+    if isinstance(result, str):
         return Response.from_message(message, text=result)
     # TODO handle attachments, other special returns
-    else:
-        return Response.from_message(message, text=str(result))
+    return Response.from_message(message, text=str(result))
 
 
 def _on_command_impl(name: str, fn: Callable[..., Any]) -> Callable[..., Any]:

lib/rollbot/injection/util.py

@@ -8,25 +8,40 @@ __all__ = [
     "Lazy",
     "MessageInjector",
     "ContextInjector",
+    "Origin",
+    "Channel",
+    "Sender",
+    "Timestamp",
+    "OriginAdmin",
+    "ChannelAdmin",
+    "Text",
+    "Attachments",
+    "CommandInjector",
 ]
 
 
-class MessageInjector(Injector[Message]):
-    async def inject(self, message: Message, context: Context) -> Message:
-        return message
-
-
-MessageInjector = MessageInjector()
+Dep = TypeVar("Dep")
 
 
-class ContextInjector(Injector[Context]):
-    async def inject(self, message: Message, context: Context) -> Context:
-        return context
+class Simple(Injector[Dep]):
+    def __init__(self, extract: Callable[[Message, Context], Dep]):
+        self.extract = extract
 
+    async def inject(self, message: Message, context: Context) -> Dep:
+        return self.extract(message, context)
 
-ContextInjector = ContextInjector()
 
-Dep = TypeVar("Dep")
+MessageInjector = Simple(lambda m, c: m)
+ContextInjector = Simple(lambda m, c: c)
+Origin = Simple(lambda m, c: m.origin_id)
+Channel = Simple(lambda m, c: m.channel_id)
+Sender = Simple(lambda m, c: m.sender_id)
+Timestamp = Simple(lambda m, c: m.timestamp)
+OriginAdmin = Simple(lambda m, c: m.origin_admin)
+ChannelAdmin = Simple(lambda m, c: m.channel_admin)
+Text = Simple(lambda m, c: m.text)
+Attachments = Simple(lambda m, c: m.attachments)
+CommandInjector = Simple(lambda m, c: m.command)
 
 
 class Lazy(InjectorWithCleanup[Callable[[], Coroutine[None, None, Dep]]]):

repl_driver.py

@@ -2,6 +2,7 @@ from datetime import datetime
 import asyncio
 
 import rollbot
+from rollbot.injection import ArgList, Lazy, Database
 
 
 class MyBot(rollbot.Rollbot[str]):
@@ -44,7 +45,7 @@ async def count_command(message, context):
 
 
 @rollbot.as_command
-async def count2(args: rollbot.ArgList, connect: rollbot.Lazy(rollbot.Database)):
+async def count2(args: ArgList, connect: Lazy(Database)):
     name = args[0] if len(args) > 0 else "main"
     db = await connect()
     await db.execute(