create a class for storing commands

Neko-Chan1 · Sep 15, 2020 · 8f70e23 · 8f70e23
1 parent c6767a0
commit 8f70e23
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 38 deletions.
diff --git a/ubot/command.py b/ubot/command.py
@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+
+from types import FunctionType
+
+
+class Command:
+    def __init__(self, func: FunctionType, args: dict):
+        self.function = func
+        self.pattern = args.get("pattern")
+        self.simple_pattern = args.get("simple_pattern", None)
+        self.raw_pattern = args.get("raw_pattern", None)
+        self.pattern_extra = args.get("pattern_extra", "")
+        self.extra = args.get("extra", None)
diff --git a/ubot/command_handler.py b/ubot/command_handler.py
@@ -19,23 +19,23 @@ def __init__(self, client, settings):
     async def handle_outgoing(self, event):
         prefix = "|".join([escape(i) for i in (self.settings.get_list("cmd_prefix") or ['.'])])
 
-        for value in self.outgoing_commands:
-            if value["simple_pattern"]:
-                pattern_match = search(self.simple_pattern_template.format(value["pattern"], value["pattern_extra"]), event.raw_text)
-            elif value["raw_pattern"]:
-                pattern_match = search(self.raw_pattern_template.format(value["pattern"] + value["pattern_extra"]), event.raw_text)
+        for command in self.outgoing_commands:
+            if command.simple_pattern:
+                pattern_match = search(self.simple_pattern_template.format(command.pattern, command.pattern_extra), event.raw_text)
+            elif command.raw_pattern:
+                pattern_match = search(self.raw_pattern_template.format(command.pattern + command.pattern_extra), event.raw_text)
             else:
-                pattern_match = search(self.pattern_template.format(prefix, value["pattern"], value["pattern_extra"]), event.raw_text)
+                pattern_match = search(self.pattern_template.format(prefix, command.pattern, command.pattern_extra), event.raw_text)
 
             if pattern_match:
                 event.pattern_match = pattern_match
                 event.args = pattern_match.groups()[-1].strip()
                 event.other_args = pattern_match.groups()[2:-1]
                 event.command = pattern_match.groups()[1]
-                event.extra = value["extra"]
+                event.extra = command.extra
 
                 try:
-                    await value["function"](event)
+                    await command.function(event)
                 except Exception as exception:
-                    await event.reply(f"`An error occurred in {value['function'].__name__}: {exception}`")
+                    await event.reply(f"`An error occurred in {command.command.__name__}: {exception}`")
                     raise exception
diff --git a/ubot/loader.py b/ubot/loader.py
@@ -6,8 +6,10 @@
 from os.path import basename, dirname, isfile
 
 from aiohttp import ClientSession
+
 from telethon.tl.types import DocumentAttributeFilename
 
+from .command import Command
 from .command_handler import CommandHandler
 
 
@@ -52,31 +54,21 @@ def reload_all_modules(self):
         return errors or None
 
     def add(self, pattern: str = None, **args):
-        pattern = args.get("pattern", pattern)
-        pattern_extra = args.get("pattern_extra", "")
-
         def decorator(func):
             if func.__module__.split(".")[-1] in self.help_dict:
                 self.help_dict[func.__module__.split(".")[-1]] += [[pattern, args.get('help', None)]]
             else:
                 self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]]
 
-            self.command_handler.outgoing_commands.append({
-                "pattern": pattern,
-                "pattern_extra": pattern_extra,
-                "function": func,
-                "simple_pattern": args.get('simple_pattern', False),
-                "raw_pattern": args.get('raw_pattern', False),
-                "extra": args.get('extra', pattern)
-            })
+            args["pattern"] = args.get("pattern", pattern)
+            self.command_handler.outgoing_commands.append(Command(func, args))
 
             return func
 
         return decorator
 
     def add_list(self, pattern: list = None, **args):
         pattern_list = args.get("pattern", pattern)
-        pattern_extra = args.get("pattern_extra", "")
 
         def decorator(func):
             for pattern in pattern_list:
@@ -85,22 +77,16 @@ def decorator(func):
                 else:
                     self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]]
 
-                self.command_handler.outgoing_commands.append({
-                    "pattern": pattern,
-                    "pattern_extra": pattern_extra,
-                    "function": func,
-                    "simple_pattern": args.get('simple_pattern', False),
-                    "raw_pattern": args.get('raw_pattern', False),
-                    "extra": args.get('extra', None)
-                })
+                this_args = args.copy()
+                this_args["pattern"] = pattern
+                self.command_handler.outgoing_commands.append(Command(func, this_args))
 
             return func
 
         return decorator
 
     def add_dict(self, pattern: dict = None, **args):
         pattern_dict = args.get("pattern", pattern)
-        pattern_extra = args.get("pattern_extra", "")
 
         def decorator(func):
             for pattern, extra in pattern_dict.items():
@@ -109,14 +95,10 @@ def decorator(func):
                 else:
                     self.help_dict[func.__module__.split(".")[-1]] = [[pattern, args.get('help', None)]]
 
-                self.command_handler.outgoing_commands.append({
-                    "pattern": pattern,
-                    "pattern_extra": pattern_extra,
-                    "function": func,
-                    "simple_pattern": args.get('simple_pattern', False),
-                    "raw_pattern": args.get('raw_pattern', False),
-                    "extra": args.get('extra', extra)
-                })
+                this_args = args.copy()
+                this_args["pattern"] = pattern
+                this_args["extra"] = args.get('extra', extra)
+                self.command_handler.outgoing_commands.append(Command(func, this_args))
 
             return func