type the event handler (dispatch)

Cog-Creators · Mar 29, 2022 · 08368ba · 08368ba
1 parent 68e4922
commit 08368ba
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/lavalink/lavalink.py b/lavalink/lavalink.py
@@ -1,5 +1,5 @@
 import asyncio
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import discord
 from discord.ext.commands import Bot
@@ -261,7 +261,9 @@ async def _handle_update(player, data: node.PositionTime, raw_data: dict):
     await player.handle_player_update(data)
 
 
-def _get_update_args(data: enums.PlayerState, raw_data: dict):
+def _get_update_args(
+    data: Union[node.PositionTime, enums.LavalinkEvents, node.Stats], raw_data: dict
+):
     guild_id = int(raw_data.get("guildId"))
 
     try:
@@ -327,7 +329,11 @@ def unregister_stats_listener(coro):
         pass
 
 
-def dispatch(op: enums.LavalinkIncomingOp, data, raw_data: dict):
+def dispatch(
+    op: enums.LavalinkIncomingOp,
+    data: Union[node.PositionTime, enums.LavalinkEvents, node.Stats],
+    raw_data: dict,
+):
     listeners = []
     args = []
     if op == enums.LavalinkIncomingOp.EVENT:

diff --git a/lavalink/node.py b/lavalink/node.py
@@ -418,12 +418,12 @@ async def _handle_op(self, op: LavalinkIncomingOp, data):
                 self.event_handler(op, event, data)
         elif op == LavalinkIncomingOp.PLAYER_UPDATE:
             state = data.get("state", {})
-            state = PositionTime(
+            position = PositionTime(
                 position=state.get("position", 0),
                 time=state.get("time", 0),
                 connected=state.get("connected", False),
             )
-            self.event_handler(op, state, data)
+            self.event_handler(op, position, data)
         elif op == LavalinkIncomingOp.STATS:
             stats = Stats(
                 memory=data.get("memory"),