From 62f08cbf4c7df3993f75e7a5cd2df9745ae6fd55 Mon Sep 17 00:00:00 2001
From: Kae Bartlett <kae@voxelfox.co.uk>
Date: Mon, 11 Sep 2023 20:25:20 +0100
Subject: [PATCH] Change autocomplete behaviour

---
 novus/ext/client/command.py | 82 ++++++++++++++++---------------------
 1 file changed, 36 insertions(+), 46 deletions(-)

diff --git a/novus/ext/client/command.py b/novus/ext/client/command.py
index 416f37da..c3f88405 100644
--- a/novus/ext/client/command.py
+++ b/novus/ext/client/command.py
@@ -42,15 +42,22 @@
             Coroutine[Any, Any, None],
         ],
     ]
-    APT = TypeVar("APT")
-    AutocompleteCallback = Callable[
-        [
-            Any,
-            n.Interaction[n.ApplicationCommandData],
-            list[n.InteractionOption],
-            str,
+    AutocompleteCallback = Union[
+        Callable[
+            [
+                Any,
+                n.Interaction[n.ApplicationCommandData],
+                dict[str, n.InteractionOption],
+            ],
+            Coroutine[Any, Any, list[n.ApplicationCommandChoice] | None],
+        ],
+        Callable[
+            [
+                Any,
+                n.Interaction[n.ApplicationCommandData],
+            ],
+            Coroutine[Any, Any, list[n.ApplicationCommandChoice] | None],
         ],
-        Coroutine[Any, Any, list[n.ApplicationCommandChoice]],
     ]
 
     LocType = dict[str, str] | dict[n.Locale, str] | n.utils.Localization | None
@@ -67,23 +74,6 @@
 log = logging.getLogger("novus.ext.client.command")
 
 
-class Autocomplete:
-    """
-    A function wrapper for autocomplete objects.
-
-    Parameters
-    ----------
-    name : str
-        The name of the option that the autocomplete is a part of.
-    callback
-        The function that acts as the autocomplete.
-    """
-
-    def __init__(self, name: str, callback: AutocompleteCallback):
-        self.name = name
-        self.callback = callback
-
-
 class Command:
     """
     A command object for Novus command handling.
@@ -145,7 +135,7 @@ def __init__(
             and " " in self.name
         )
         self.owner: Any = None
-        self.autocompletes: dict[str, Autocomplete] = {}
+        self._autocomplete = None
 
         # Make sure our callback and app command have similar options
         if self.type == n.ApplicationCommandType.chat_input:
@@ -289,7 +279,7 @@ async def run(
     async def __call__(self, *args, **kwargs) -> Any:
         return await self.callback(self.owner, *args, **kwargs)
 
-    def autocomplete(self, name: str) -> Callable[[AutocompleteCallback], AutocompleteCallback]:
+    def autocomplete(self, func: AutocompleteCallback) -> None:
         """
         Add an autocomplete to this command.
 
@@ -299,11 +289,7 @@ def autocomplete(self, name: str) -> Callable[[AutocompleteCallback], Autocomple
             The name of the option that should be autocompleted.
         """
 
-        def wrapper(func: AutocompleteCallback):
-            autocomplete = Autocomplete(name, func)
-            self.autocompletes[name] = autocomplete
-            return func
-        return wrapper
+        self._autocomplete = func
 
     async def run_autocomplete(
             self,
@@ -319,22 +305,26 @@ async def run_autocomplete(
             The interaction that needs completing.
         """
 
-        autocomplete: Autocomplete | None = None
-        current_value: str | int | float | bool | None = None
+        if self._autocomplete is None:
+            return
         options = options or interaction.data.options
-        for i in options:
-            if i.focused:
-                autocomplete = self.autocompletes.get(i.name)
-                current_value = i.value
-                break
-        if autocomplete is None:
+        if hasattr(self._autocomplete, "_param_count"):
+            param_count = self._autocomplete._param_count
+        else:
+            param_count = self._autocomplete._param_count = len(inspect.signature(self._autocomplete).parameters)
+        if param_count == 3:
+            data = await self._autocomplete.callback(
+                self.owner,
+                interaction,
+                {i.name: i for i in options},
+            )
+        else:
+            data = await self._autocomplete.callback(
+                self.owner,
+                interaction,
+            )
+        if data is None:
             return await interaction.send_autocomplete([])
-        data = await autocomplete.callback(
-            self.owner,
-            interaction,
-            options,
-            current_value,  # pyright: ignore
-        )
         return await interaction.send_autocomplete(data)