feat: component method decorators

TheWii · Jun 17, 2022 · af94217 · af94217
1 parent 282d825
commit af94217
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/integrity/components.py b/integrity/components.py
@@ -101,6 +101,17 @@ def run(self, event_name: str):
     def setmethod(self, name: str, function: Callable[["Component"], Any]):
         self.methods[name] = partial(function, self)
 
+    def method(self, arg):
+        def decorator(func):
+            self.setmethod(arg, func)
+            return func
+
+        if isinstance(arg, str):
+            return decorator
+        if callable(arg):
+            self.setmethod(arg.__name__, arg)
+            return arg
+
     def __getattr__(self, key: str):
         return self.methods[key]