From c874017a6b3cbe3e7f59af361a1b1f0ca4558b8f Mon Sep 17 00:00:00 2001
From: Zomatree <angelokontaxis@hotmail.com>
Date: Sat, 13 Nov 2021 23:52:23 +0000
Subject: [PATCH] add more checks

---
 revolt/ext/commands/checks.py | 41 ++++++++++++++++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/revolt/ext/commands/checks.py b/revolt/ext/commands/checks.py
index de66f65..d40c38a 100755
--- a/revolt/ext/commands/checks.py
+++ b/revolt/ext/commands/checks.py
@@ -1,14 +1,28 @@
 from typing import Callable, Coroutine, Any, TypeVar, Union
+
 from .command import Command
 from .context import Context
+from .errors import NotBotOwner, NotServerOwner, ServerOnly
 
-__all__ = ("check", "Check")
+__all__ = ("check", "Check", "is_bot_owner", "is_server_owner")
 
 Check = Callable[[Context], Union[Any, Coroutine[Any, Any, Any]]]
 
 T = TypeVar("T", Callable[..., Any], Command)
 
 def check(check: Check):
+    """A decorator for adding command checks
+
+    Parameters
+    -----------
+    check: Callable[[Context], Union[Any, Coroutine[Any, Any, Any]]]
+        The function to be called, must take one parameter, context and optionally be a coroutine
+
+    Returns
+    --------
+    Any
+        The value denoating whether the check should pass or fail
+    """
     def inner(func: T) -> T:
         if isinstance(func, Command):
             func.checks.append(check)
@@ -20,3 +34,28 @@ def check(check: Check):
         return func
 
     return inner
+
+def is_bot_owner():
+    """A command check for limiting the command to only the bot's owner"""
+    @check
+    def inner(context: Context):
+        if context.author.id == context.client.user.owner_id:
+            return True
+
+        raise NotBotOwner
+
+    return inner
+
+def is_server_owner():
+    """A command check for limiting the command to only a server's owner"""
+    @check
+    def inner(context: Context):
+        if not context.server:
+            raise ServerOnly
+
+        if context.author.id == context.server.owner_id:
+            return True
+
+        raise NotServerOwner
+
+    return inner