Browse Source

Provide OptArg injection

Kirk Trombley 4 years ago
parent
commit
9d375a2395
1 changed files with 28 additions and 1 deletions
  1. 28 1
      lib/rollbot/plugins/injection.py

+ 28 - 1
lib/rollbot/plugins/injection.py

@@ -24,7 +24,7 @@ class Arg(ArgConverter):
         super().__init__(self._convert)
         self.index = index
         self.conversion = conversion
-        self.fail_msg = fail_msg
+        self.fail_msg = fail_msg or "Invalid argument: {}"
     
     def _convert(self, cmd, db, msg):
         try:
@@ -37,12 +37,39 @@ class Arg(ArgConverter):
             RollbotFailure.INVALID_ARGUMENTS.with_reason(self.fail_msg.format(arg)).raise_exc()
 
 
+class OptArg(ArgConverter):
+    def __init__(self, index, conversion=str, default=None, fail_msg=None):
+        super().__init__(self._convert)
+        self.index = index
+        self.conversion = conversion
+        self.default = default
+        self.fail_msg = fail_msg or "Invalid argument: {}"
+    
+    def _convert(self, cmd, db, msg):
+        if msg.raw_args is None:
+            return self.default
+        try:
+            arg = msg.arg_list()[self.index]
+        except IndexError:
+            return self.default
+        try:
+            return self.conversion(arg)
+        except ValueError:
+            RollbotFailure.INVALID_ARGUMENTS.with_reason(self.fail_msg.format(arg)).raise_exc()
+
+
 class _SubcArg(Arg):
     def _convert(self, cmd, db, msg):
         return super()._convert(self, cmd, db, RollbotMessage.from_subcommand(msg))
 
 
+class _SubcOptArg(OptArg):
+    def _convert(self, cmd, db, msg):
+        return super()._convert(self, cmd, db, RollbotMessage.from_subcommand(msg))
+
+
 Subcommand.Arg = _SubcArg
+Subcommand.OptArg = _SubcOptArg
 
 
 class Config(ArgConverter):