feat: add charge detect and new tray icon

Author: Maasea

src/controller.py

@@ -0,0 +1,85 @@
+import time
+import usb.core
+import usb.util
+from usb.backend import libusb1
+from storage import RebtStorage, DEVICE
+
+
+# You can learn more info from https://github.com/hsutungyu/razer-mouse-battery-windows/blob/main/mamba.pyw
+class Rebt:
+
+    def __init__(self) -> None:
+        self.config = RebtStorage()
+        self.findAll = True
+        self.hasDefault = self.config.isExist()
+        self.firstRun = not self.hasDefault
+        self.backend = libusb1.get_backend(find_library=lambda x: "../libusb-1.0.dll")
+        self.mouse = self.get_mouse()
+
+    def map_mouse(self):
+        mouses = usb.core.find(find_all=True, idVendor=0x1532, backend=self.backend)
+        for mouse in mouses:
+            for name, device in DEVICE.items():
+                knowProduct = int(device["usbId"], 16)
+                if knowProduct == mouse.idProduct:
+                    self.config.generate(name, device["usbId"], device["tranId"])
+                    self.hasDefault = True
+                    return mouse
+
+    def get_mouse(self):
+        mouse = None
+        if self.hasDefault:
+            idProduct = self.config.getDefault("usbId")
+            mouse = usb.core.find(idVendor=0x1532, idProduct=idProduct, backend=self.backend)
+        if not mouse and self.findAll:
+            mouse = self.map_mouse()
+        if not mouse:
+            raise RuntimeError(f"Mouse cannot be found.")
+
+        return mouse
+
+    def generate_msg(self, transaction_id, command_class, command_id, data_size):
+        msgs = [0x00, transaction_id, 0x00, 0x00, 0x00, command_class, command_id, data_size]
+        msg = bytes(msgs)
+        crc = 0
+        for i in msg[2:]:
+            crc ^= i
+        msg += bytes(80)
+        msg += bytes([crc, 0])
+        return msg
+
+    def battery_msg(self):
+        return self.generate_msg(self.config.getDefault("tranId"), 0x02, 0x07, 0x80)
+
+    def charge_msg(self):
+        return self.generate_msg(self.config.getDefault("tranId"), 0x02, 0x07, 0x84)
+
+    def send_msg(self, msg):
+        self.mouse = self.get_mouse()
+        usb.util.claim_interface(self.mouse, 0)
+        self.mouse.set_configuration()
+        self.mouse.ctrl_transfer(bmRequestType=0x21,
+                                 bRequest=0x09,
+                                 wValue=0x300,
+                                 data_or_wLength=msg,
+                                 wIndex=0x00)
+        time.sleep(0.1)
+        result = self.mouse.ctrl_transfer(bmRequestType=0xa1,
+                                          bRequest=0x01,
+                                          wValue=0x300,
+                                          data_or_wLength=90,
+                                          wIndex=0x00)
+        usb.util.dispose_resources(self.mouse)
+        usb.util.release_interface(self.mouse, 0)
+        return result[9]
+
+    def get_battery(self):
+        battery_msg = self.battery_msg()
+        battery = self.send_msg(battery_msg)
+        if battery == 0: return -1
+        return int(battery / 255 * 100)
+
+    def is_charging(self):
+        # 1 charging, 0 not
+        charge_msg = self.charge_msg()
+        return self.send_msg(charge_msg)

src/drawtray.py

@@ -0,0 +1,82 @@
+from PySide2.QtCore import Qt, QRect, QRectF, QSize, QPointF
+from PySide2.QtGui import QPixmap, QPainter, QFont, QColor, QPen, QFontMetrics
+from PySide2.QtWidgets import QApplication
+
+
+class TrayIcon:
+
+    def __init__(self):
+        self.bgWidth = 128
+        self.bgHeight = 128
+        self.pixmap = QPixmap(self.bgWidth, self.bgHeight)
+        self.pixmap.fill(Qt.transparent)
+        self.painter = QPainter(self.pixmap)
+
+    def centerWidth(self, width):
+        return (self.bgWidth - width) >> 1
+
+    def centerHeight(self, height):
+        return (self.bgHeight - height) >> 1
+
+    def drawRoundedRect(self, margin, radius, battery=1.0):
+        marginX = margin[0]
+        marginY = margin[1]
+        bRectWidth = self.bgWidth - marginX * 2
+        bRectHeight = self.bgHeight - marginY * 2
+        brt = QRectF(self.centerWidth(bRectWidth), self.centerHeight(bRectHeight), bRectWidth * battery, bRectHeight)
+        self.painter.drawRoundedRect(brt, radius, radius)
+        return brt
+
+    def drawChord(self, brt):
+        width = height = 15
+        cr = QRectF(brt.x() + brt.width() - 5, self.centerHeight(height), width, height)
+        self.painter.drawChord(cr, 90 * 16, -180 * 16)
+
+    def drawChargeIcon(self, r=20):
+        p = QRectF(self.pixmap.width() - r, 0, r, r)
+        self.painter.setPen(Qt.red)
+        self.painter.setBrush(Qt.red)
+        self.painter.drawChord(p, 0, 360 * 16)
+
+    def drawNumIcon(self, battery, isCharge):
+        if isCharge: self.drawChargeIcon(26)
+        ft = QFont("Microsoft YaHei", 48)
+        ft_rect = QFontMetrics(ft).boundingRect(str(battery))
+        scale = ft_rect.width() / self.pixmap.width()
+        if scale > 1:
+            ft.setPointSize(36)
+        elif scale < 0.5:
+            ft.setPointSize(62)
+        self.painter.setFont(ft)
+        self.painter.setPen(Qt.black)
+        self.painter.setRenderHint(QPainter.HighQualityAntialiasing)
+        self.painter.drawText(self.pixmap.rect(), Qt.AlignCenter, str(battery))
+
+        self.painter.end()
+
+    def drawRectIcon(self, battery, isCharge):
+        self.painter.setBrush(Qt.black)
+        self.painter.setPen(Qt.transparent)
+        brt = self.drawRoundedRect((16, 42), 6)
+
+        self.painter.setCompositionMode(QPainter.CompositionMode_SourceOut)
+        percentage = battery / 100
+        mask_x = brt.x() + percentage * brt.width()
+        mask_width = brt.width() * (1 - percentage)
+        maskBrt = QRectF(mask_x, brt.y(), mask_width, brt.height())
+        self.painter.fillRect(maskBrt, Qt.transparent)
+
+        # big rect
+        self.painter.setCompositionMode(QPainter.CompositionMode_DestinationOver)
+        self.painter.setPen(QPen(Qt.black, 9))
+        self.painter.setBrush(Qt.transparent)
+        brt_ = self.drawRoundedRect((6, 30), 16)
+        if isCharge: self.drawChargeIcon()
+        self.painter.end()
+
+    def draw(self, battery, isCharge, style=0):
+        if style == 0:
+            self.drawNumIcon(battery, isCharge)
+        else:
+            self.drawRectIcon(battery, isCharge)
+        return self.pixmap