refactor: optimize how autodefer works

interactions/client/models/utils.py

@@ -1,12 +1,23 @@
 from asyncio import Task, get_running_loop, sleep
 from functools import wraps
-from typing import TYPE_CHECKING, Awaitable, Callable, Iterable, List, Optional, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Awaitable,
+    Callable,
+    Coroutine,
+    Iterable,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+)
 
 from ...api.error import LibraryException
 from .component import ActionRow, Button, SelectMenu
 
 if TYPE_CHECKING:
     from ..context import CommandContext
+    from ..bot import Extension
 
 __all__ = ("autodefer", "spread_to_rows", "search_iterable")
 
@@ -41,16 +52,24 @@ async def command(ctx):
     :rtype:
     """
 
-    def decorator(coro: Callable[..., Awaitable]) -> Callable[..., Awaitable]:
+    def decorator(coro: Callable[..., Union[Awaitable, Coroutine]]) -> Callable[..., Awaitable]:
         from ..context import ComponentContext
 
         @wraps(coro)
-        async def deferring_func(ctx: Union["CommandContext", "ComponentContext"], *args, **kwargs):
+        async def deferring_func(ctx: Union["CommandContext", "ComponentContext", "Extension"], *args, **kwargs):
             try:
                 loop = get_running_loop()
             except RuntimeError as e:
                 raise RuntimeError("No running event loop detected!") from e
-            task: Task = loop.create_task(coro(ctx, *args, **kwargs))
+
+            if isinstance(args[0], (ComponentContext, CommandContext)):
+                self = ctx
+                ctx = list(args).pop(0)
+
+                task: Task = loop.create_task(coro(self, ctx, *args, **kwargs))
+
+            else:
+                task: Task = loop.create_task(coro(ctx, *args, **kwargs))
 
             await sleep(delay)