Browse Source

First pass at subcommand logic, pretty untested

Kirk Trombley 4 years ago
parent
commit
33f4ed136d

+ 3 - 6
lib/rollbot/injection/args.py

@@ -5,7 +5,7 @@ import shlex
 
 from ..types import Message, Context
 from ..failure import RollbotFailure
-from .base import Injector
+from .base import Injector, Simple
 
 __all__ = [
     "Args",
@@ -16,9 +16,7 @@ __all__ = [
 ]
 
 
-class ArgsInjector(Injector[str]):
-    async def inject(self, message: Message, context: Context) -> str:
-        return message.command.args
+Args = Simple[str](lambda m, c: m.command.args)
 
 
 class ArgListSplitOn(Injector[list[str]]):
@@ -43,10 +41,9 @@ class ArgParse(Injector[Namespace]):
         self.parser = parser
 
     async def inject(self, message: Message, context: Context) -> Namespace:
-        return self.parser.parse_args(shlex.split(message.text))
+        return self.parser.parse_args(shlex.split(message.command.args))
 
 
-Args = ArgsInjector()
 ArgList = ArgListSplitOn()
 
 ArgType = TypeVar("ArgType")

+ 9 - 0
lib/rollbot/injection/base.py

@@ -1,3 +1,4 @@
+from collections.abc import Callable
 from typing import Generic, TypeVar, Any
 from contextlib import asynccontextmanager
 
@@ -16,6 +17,14 @@ class InjectorWithCleanup(Injector[Dep]):
         raise NotImplementedError
 
 
+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)
+
+
 @asynccontextmanager
 async def inject_all(injectors: list[Injector[Any]], message: Message, context: Context):
     try:

+ 19 - 0
lib/rollbot/injection/subcommand.py

@@ -0,0 +1,19 @@
+from ..types import Message, Context, Command
+from .base import Injector
+
+__all__ = [
+    "Subcommand",
+]
+
+
+class SubcommandInjector(Injector[Command]):
+    async def inject(self, message: Message, context: Context) -> Command:
+        subc = message.command.cache.get(SubcommandInjector.__name__, None)
+        if subc is None:
+            subc = message.command.get_subcommand()
+            message.command.cache[SubcommandInjector.__name__] = subc
+        return subc
+
+
+Subcommand = SubcommandInjector()
+# TODO Subcommand.Arg, .Args, .ArgList, .ArgListSplitOn, .ArgParse

+ 4 - 13
lib/rollbot/injection/util.py

@@ -5,7 +5,7 @@ from datetime import datetime
 from aiohttp import ClientSession
 
 from ..types import Message, Context, Attachment, Command
-from .base import Injector, InjectorWithCleanup
+from .base import Injector, InjectorWithCleanup, Simple
 
 __all__ = [
     "Lazy",
@@ -25,18 +25,6 @@ __all__ = [
     "Config",
 ]
 
-
-Dep = TypeVar("Dep")
-
-
-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)
-
-
 MessageInjector = Simple[Message](lambda m, c: m)
 ContextInjector = Simple[Context](lambda m, c: c)
 Origin = Simple[str](lambda m, c: m.origin_id)
@@ -60,6 +48,9 @@ class Config(Injector[Any]):
         return context.config(self.key)
 
 
+Dep = TypeVar("Dep")
+
+
 class Lazy(InjectorWithCleanup[Callable[[], Coroutine[None, None, Dep]]]):
     def __init__(self, deferred: Injector[Dep]):
         self.deferred = deferred

+ 5 - 0
lib/rollbot/types.py

@@ -68,6 +68,11 @@ class Command:
             args=parts[1] if len(parts) > 1 else "",
         )
 
+    def get_subcommand(self, inherit_bang=True) -> Command:
+        if inherit_bang and not self.args.startswith(self.bang):
+            return Command.from_text(self.bang + self.args)
+        return Command.from_text(self.args)
+
 
 @dataclass
 class Response: