feat: pylance basic type

fix: websockets.ConnectionClosed no await
style: isort

Author: phillychi3

.vscode/settings.json

@@ -3,5 +3,6 @@
         "tests"
     ],
     "python.testing.unittestEnabled": false,
-    "python.testing.pytestEnabled": true
+    "python.testing.pytestEnabled": true,
+    "python.analysis.typeCheckingMode": "basic"
 }

dglabv3/__init__.py

@@ -1,3 +1,3 @@
 from .dglab import dglabv3  # noqa: F401
-from .dtype import Channel, StrengthType, Button, Strength  # noqa: F401
-from .waves import PULSES, ALL_PULSES, Pulse  # noqa: F401
+from .dtype import Button, Channel, Strength, StrengthType  # noqa: F401
+from .waves import ALL_PULSES, PULSES, Pulse  # noqa: F401

dglabv3/dglab.py

@@ -1,16 +1,19 @@
+import asyncio
+import io
 import json
 import logging
-import qrcode
-import io
-import asyncio
 from threading import Event
-from websockets.asyncio.client import connect as ws_connect
+from typing import Optional
+
+import qrcode
 import websockets
+from websockets.asyncio.client import connect as ws_connect
 
-from dglabv3.dtype import Button, Channel, StrengthType, StrengthMode, MessageType, ChannelStrength, Strength
-from dglabv3.wsmessage import WSMessage, WStype
+from dglabv3.dtype import (Button, Channel, ChannelStrength, MessageType,
+                           Strength, StrengthMode, StrengthType)
 from dglabv3.event import EventEmitter
 from dglabv3.music_to_wave import convert_audio_to_v3_protocol
+from dglabv3.wsmessage import WSMessage, WStype
 
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger("dglabv3")
@@ -58,7 +61,7 @@ def is_connected(self) -> bool:
         """
         是否連接到WebSocket
         """
-        return self.client and self.client.sock and self.client.sock.connected
+        return self.client is not None
 
     def is_linked_to_app(self) -> bool:
         """
@@ -110,6 +113,9 @@ async def connect(self) -> None:
 
     async def _listen(self):
         try:
+            if self.client is None:
+                logger.error("WebSocket client is None")
+                return
             async for message in self.client:
                 await self._handle_message(message)
         except websockets.ConnectionClosed:
@@ -118,7 +124,7 @@ async def _listen(self):
             logger.error(f"WebSocket error: {e}")
             raise ConnectionError("WebSocket error")
 
-    def generate_qrcode(self) -> io.BytesIO:
+    def generate_qrcode(self) -> Optional[io.BytesIO]:
         """
         生成QR code圖片
         """
@@ -130,11 +136,11 @@ def generate_qrcode(self) -> io.BytesIO:
         qr.make(fit=True)
         img = qr.make_image(fill_color="black", back_color="white")
         saveimg = io.BytesIO()
-        img.save(saveimg, format="PNG")
+        img.save(saveimg, format="PNG")  # type: ignore
         saveimg.seek(0)
         return saveimg
 
-    def generate_qrcode_text(self) -> str:
+    def generate_qrcode_text(self) -> Optional[str]:
         """
         生成QR code文字
         """
@@ -174,7 +180,7 @@ async def _heartbeat(self):
 
         except websockets.ConnectionClosed:
             logger.info("WebSocket connection closed")
-            self.close()
+            await self.close()
         except Exception as e:
             logger.error(f"Heartbeat error: {e}")
 
@@ -184,7 +190,7 @@ def _start_heartbeat(self):
             self._heartbeat_task.cancel()
         self._heartbeat_task = asyncio.create_task(self._heartbeat())
 
-    async def _handle_message(self, data: str):
+    async def _handle_message(self, data: websockets.Data):
         try:
             message = json.loads(data)
             WSmsg = WSMessage(message)
@@ -195,14 +201,17 @@ async def _handle_message(self, data: str):
                 self._bind_event.set()
 
             elif WSmsg.type == WStype.MSG:
-                if WSmsg.msg.startswith("feedback"):
-                    button = WSmsg.feedback()
-                    await self._dispatch_button(button)
-                elif WSmsg.msg.startswith("strength"):
-                    self.strength.set_strength(WSmsg.strength())
-                    await self._dispatch_strength(WSmsg.strength())
+                if WSmsg.msg is not None:
+                    if WSmsg.msg.startswith("feedback"):
+                        button = WSmsg.feedback()
+                        await self._dispatch_button(button)
+                    elif WSmsg.msg.startswith("strength"):
+                        self.strength.set_strength(WSmsg.strength())
+                        await self._dispatch_strength(WSmsg.strength())
+                    else:
+                        logger.warning(f"Unknown message type: {WSmsg.msg}")
                 else:
-                    logger.warning(f"Unknown message type: {WSmsg.msg}")
+                    logger.warning("Received message with None content")
 
             logger.debug(f"Received message: {message}")
         except Exception as e:
@@ -263,18 +272,19 @@ async def send_wave_message(self, wave: list[list[list[int]]], time: int = 10, c
         time: 30\n
         channel: Channel.A
         """
-        if channel == 1:
-            channel = "A"
-        elif channel == 2:
-            channel = "B"
-        elif channel == 3:
-            channel = "BOTH"
+        channel_str = ""
+        if channel == Channel.A:
+            channel_str = "A"
+        elif channel == Channel.B:
+            channel_str = "B"
+        elif channel == Channel.BOTH:
+            channel_str = "BOTH"
 
-        def _create_wave_message(channel: str, wave, time: int) -> dict:
+        def _create_wave_message(ch_str: str, wave, time: int) -> dict:
             return {
                 "type": MessageType.CLIENT_MSG,
-                "channel": channel,
-                "message": f"{channel}:{json.dumps(self._wave2hex(wave))}",
+                "channel": ch_str,
+                "message": f"{ch_str}:{json.dumps(self._wave2hex(wave))}",
                 "time": time,
             }
 
@@ -283,12 +293,12 @@ def _create_wave_message(channel: str, wave, time: int) -> dict:
         # message2 : B通道波形数据(16进制HEX数组json,具体见上面的协议说明)
         # time1 : A通道波形数据持续发送时长
         # time2 : B通道波形数据持续发送时长
-        if channel == "BOTH":
+        if channel_str == "BOTH":
             for ch in ["A", "B"]:
                 message = _create_wave_message(ch, wave, time)
                 await self._send_message(message)
         else:
-            message = _create_wave_message(channel, wave, time)
+            message = _create_wave_message(channel_str, wave, time)
             await self._send_message(message)
 
     async def clear_wave(self, channel: Channel):

dglabv3/dtype.py

@@ -1,6 +1,6 @@
+from dataclasses import dataclass, field
 from enum import Enum, IntEnum, StrEnum
 from typing import Final
-from dataclasses import dataclass, field
 
 __all__ = ["ChannelStrength", "StrengthType", "StrengthMode", "MessageType", "Channel"]
 

dglabv3/event.py

@@ -1,13 +1,12 @@
 import asyncio
-from typing import Callable, Any, Dict, List
 import functools
 import logging
-
+from typing import Any, Callable, Dict, List, Optional
 
 logger = logging.getLogger("dglabv3.event")
 
 
-def event(name: str = None):
+def event(name: Optional[str] = None):
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         @functools.wraps(func)
         async def async_wrapper(*args: Any, **kwargs: Any) -> Any:

dglabv3/music_to_wave.py

@@ -1,7 +1,8 @@
-import librosa
 import math
-import numpy as np
+
+import librosa
 import matplotlib.pyplot as plt
+import numpy as np
 
 
 def convert_audio_to_v3_protocol(mp3_file_path: str) -> list:

dglabv3/wsmessage.py

@@ -1,5 +1,6 @@
 from enum import Enum
 from typing import Optional
+
 from dglabv3.dtype import Button, Strength
 
 
@@ -11,13 +12,10 @@ class WStype(Enum):
     ERROR = "error"
 
 
-
-
-
 class WSMessage:
     def __init__(self, data: dict):
         self.type: WStype = WStype(data.get("type"))
-        self.msg: Optional[dict] = data.get("message", None)
+        self.msg: Optional[str] = data.get("message", None)  # 將型別從 dict 改為 str
         self.targetID: Optional[str] = data.get("targetId", None)
         self.clientID: Optional[str] = data.get("clientId", None)
 
@@ -34,8 +32,12 @@ def to_dict(self) -> dict:
         }
 
     def feedback(self) -> Button:
+        if self.msg is None:
+            raise ValueError("Message is None, cannot get feedback")
         return Button(self.msg.split("-")[1])
 
     def strength(self) -> Strength:
+        if self.msg is None:
+            raise ValueError("Message is None, cannot get strength")
         splitmsg = self.msg.split("-")[1].split("+")
         return Strength(A=int(splitmsg[0]), B=int(splitmsg[1]), MAXA=int(splitmsg[2]), MAXB=int(splitmsg[3]))

example/connect.py

@@ -1,12 +1,11 @@
 import asyncio
-from PIL import Image
-import sys
 import os
+import sys
 
-sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from dglabv3 import dglabv3
-from dglabv3 import Channel, StrengthType, PULSES
+from PIL import Image
 
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from dglabv3 import PULSES, Channel, StrengthType, dglabv3
 
 client = dglabv3()
 
@@ -15,6 +14,9 @@ async def run():
     try:
         await client.connect_and_wait()
         qrcode = client.generate_qrcode()
+        if qrcode is None:
+            print("Failed to generate QR code.")
+            return
         ig = Image.open(qrcode)
         ig.show()
         await client.wait_for_app_connect()

example/discordbot.py

@@ -1,8 +1,10 @@
+import random
+
 import discord
 from discord import app_commands
-import random
 from discord.ext import commands
-from dglabv3 import dglabv3, Channel, Pulse, ALL_PULSES, PULSES, Strength
+
+from dglabv3 import ALL_PULSES, PULSES, Channel, Pulse, Strength, dglabv3
 
 
 class dglab:

example/eventtest.py

@@ -1,13 +1,11 @@
 import asyncio
-from PIL import Image
-import sys
 import os
+import sys
 
+from PIL import Image
 
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from dglabv3 import dglabv3
-from dglabv3 import Channel, StrengthType, PULSES, Strength, Button
-
+from dglabv3 import PULSES, Button, Channel, Strength, StrengthType, dglabv3
 
 client = dglabv3()
 
@@ -26,6 +24,9 @@ async def run():
     try:
         await client.connect_and_wait()
         qrcode = client.generate_qrcode()
+        if qrcode is None:
+            print("Failed to generate QR code.")
+            return
         ig = Image.open(qrcode)
         ig.show()