makcu 0.1.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

makcu/connection.py

@@ -1,271 +1,486 @@
 import serial
 import threading
 import time
+import struct
+from typing import Optional, Dict, Callable, List, Any, Tuple, Union
 from serial.tools import list_ports
+from dataclasses import dataclass, field
+from collections import deque
+from concurrent.futures import Future
+import logging
+import asyncio
+import re
 from .errors import MakcuConnectionError, MakcuTimeoutError
 from .enums import MouseButton
 
-class SerialTransport:
-    baud_change_command = bytearray([0xDE, 0xAD, 0x05, 0x00, 0xA5, 0x00, 0x09, 0x3D, 0x00])
-
-    button_map = {
-        0: 'left',
-        1: 'right',
-        2: 'middle',
-        3: 'mouse4',
-        4: 'mouse5'
-    }
+logger = logging.getLogger(__name__)
+
+@dataclass
+class PendingCommand:
+    """Tracks a command waiting for response"""
+    command_id: int
+    command: str
+    future: Future
+    timestamp: float
+    expect_response: bool = True
+    timeout: float = 0.1  # Reduced from 1.0
+
+@dataclass
+class ParsedResponse:
+    """Parsed response from device"""
+    command_id: Optional[int]
+    content: str
+    is_button_data: bool = False
+    button_mask: Optional[int] = None
 
-    def __init__(self, fallback, debug=False, send_init=True):        
+class SerialTransport:
+    """Ultra-optimized serial transport for gaming performance"""
+    
+    BAUD_CHANGE_COMMAND = bytearray([0xDE, 0xAD, 0x05, 0x00, 0xA5, 0x00, 0x09, 0x3D, 0x00])
+    DEFAULT_TIMEOUT = 0.1  # Reduced from 1.0 for gaming
+    MAX_RECONNECT_ATTEMPTS = 3  # Reduced from 5
+    RECONNECT_DELAY = 0.1  # Reduced from 0.5
+    
+    # Pre-computed button maps for faster lookups
+    BUTTON_MAP = (
+        'left', 'right', 'middle', 'mouse4', 'mouse5'
+    )
+    
+    BUTTON_ENUM_MAP = (
+        MouseButton.LEFT,
+        MouseButton.RIGHT,
+        MouseButton.MIDDLE,
+        MouseButton.MOUSE4,
+        MouseButton.MOUSE5,
+    )
+
+    def __init__(self, fallback: str = "", debug: bool = False, 
+                 send_init: bool = True, auto_reconnect: bool = True, 
+                 override_port: bool = False) -> None:
+        # Basic config
         self._fallback_com_port = fallback
-        self._log_messages = []
         self.debug = debug
         self.send_init = send_init
-        self._button_callback = None
-        self._last_mask = 0
-        self._lock = threading.Lock()
+        self.auto_reconnect = auto_reconnect
+        self.override_port = override_port
+        
+        # Connection state
         self._is_connected = False
-        self._stop_event = threading.Event()
-        self._listener_thread = None
-        self._button_states = {btn: False for btn in self.button_map.values()}
-        self._last_callback_time = {bit: 0 for bit in self.button_map}
-        self._pause_listener = False
-
-        self._button_enum_map = {
-            0: MouseButton.LEFT,
-            1: MouseButton.RIGHT,
-            2: MouseButton.MIDDLE,
-            3: MouseButton.MOUSE4,
-            4: MouseButton.MOUSE5,
-        }
-
-        self.port = self.find_com_port()
-        if not self.port:
-            raise MakcuConnectionError("Makcu device not found. Please specify a port explicitly.")
-
+        self._reconnect_attempts = 0
+        self.port: Optional[str] = None
         self.baudrate = 115200
-        self.serial = None
-        self._current_baud = None
-
-
-    def receive_response(self, max_bytes=1024, max_lines=3, sent_command: str = "") -> str:
-        lines = []
-        try:
-            for _ in range(max_lines):
-                line = self.serial.readline(max_bytes)
-                if not line:
-                    break
-                decoded = line.decode(errors="ignore").strip()
-                if decoded:
-                    lines.append(decoded)
-        except Exception as e:
-            print(f"[RECV ERROR] {e}")
-            return ""
-
-        command_clean = sent_command.strip()
-        if lines:
-            lines.pop(-1)
-        if command_clean in lines and len(lines) > 1:
-            lines.remove(command_clean)
-        return "\n".join(lines)
-
-    def set_button_callback(self, callback):
-        self._button_callback = callback
-
-    def _log(self, message):
-        timestamp = time.strftime("%H:%M:%S")
-        entry = f"[{timestamp}] {message}"
-        self._log_messages.append(entry)
-        if len(self._log_messages) > 20:
-            self._log_messages.pop(0)
-        print(entry, flush=True)
-
-    def find_com_port(self):
-        self._log("Searching for CH343 device...")
-
+        self.serial: Optional[serial.Serial] = None
+        self._current_baud: Optional[int] = None
+        
+        # Command tracking with pre-allocated buffer
+        self._command_counter = 0
+        self._pending_commands: Dict[int, PendingCommand] = {}
+        self._command_lock = threading.Lock()
+        
+        # Response parsing with optimized buffer
+        self._parse_buffer = bytearray(1024)  # Pre-allocate
+        self._buffer_pos = 0
+        self._response_queue = deque(maxlen=100)  # Limit queue size
+        
+        # Button state with bitwise operations
+        self._button_callback: Optional[Callable[[MouseButton, bool], None]] = None
+        self._last_button_mask = 0
+        self._button_states = 0  # Use single int instead of dict
+        
+        # Threading
+        self._stop_event = threading.Event()
+        self._listener_thread: Optional[threading.Thread] = None
+        
+        # Logging optimization
+        self._log_messages: deque = deque(maxlen=100)
+        
+        # Cache for frequently used data
+        self._ascii_decode_table = bytes(range(128))  # ASCII lookup table
+
+    def _log(self, message: str, level: str = "INFO") -> None:
+        """Optimized logging - only format when needed"""
+        if not self.debug and level == "DEBUG":
+            return
+            
+        if self.debug:
+            # Use faster time formatting
+            timestamp = f"{time.time():.3f}"
+            entry = f"[{timestamp}] [{level}] {message}"
+            self._log_messages.append(entry)
+            print(entry, flush=True)
+
+    def _generate_command_id(self) -> int:
+        """Generate unique command ID - optimized with no lock for single-threaded access"""
+        self._command_counter = (self._command_counter + 1) & 0x2710  # Faster than % 10000
+        return self._command_counter
+
+    def find_com_port(self) -> Optional[str]:
+        """Optimized port finding with caching"""
+        if self.override_port:
+            return self._fallback_com_port
+            
+        # Cache the VID:PID string
+        target_hwid = "VID:PID=1A86:55D3"
+        
         for port in list_ports.comports():
-            if "VID:PID=1A86:55D3" in port.hwid.upper():
+            if target_hwid in port.hwid.upper():
                 self._log(f"Device found: {port.device}")
                 return port.device
-
+        
         if self._fallback_com_port:
-            self._log(f"Device not found. Falling back to specified port: {self._fallback_com_port}")
+            self._log(f"Using fallback: {self._fallback_com_port}")
             return self._fallback_com_port
-        else:
-            self._log("Fallback port not specified or invalid.")
-            return None
+        
+        return None
+
+    def _parse_response_line(self, line: bytes) -> ParsedResponse:
+        """Optimized parsing - work with bytes directly"""
+        # Skip decode for simple checks
+        if line.startswith(b'>>> '):
+            content = line[4:].decode('ascii', 'ignore').strip()
+            return ParsedResponse(None, content, False)
+        
+        content = line.decode('ascii', 'ignore').strip()
+        return ParsedResponse(None, content, False)
+
+    def _handle_button_data(self, byte_val: int) -> None:
+        """Optimized button handling with bitwise operations"""
+        if byte_val == self._last_button_mask:
+            return
+            
+        changed_bits = byte_val ^ self._last_button_mask
+        
+        # Use bitwise operations instead of dict lookups
+        for bit in range(5):  # Only check 5 buttons
+            if changed_bits & (1 << bit):
+                is_pressed = bool(byte_val & (1 << bit))
+                
+                # Update button state in single int
+                if is_pressed:
+                    self._button_states |= (1 << bit)
+                else:
+                    self._button_states &= ~(1 << bit)
+                
+                if self._button_callback and bit < len(self.BUTTON_ENUM_MAP):
+                    try:
+                        self._button_callback(self.BUTTON_ENUM_MAP[bit], is_pressed)
+                    except Exception:
+                        pass  # Silently ignore callback errors for speed
+        
+        self._last_button_mask = byte_val
+
+    def _process_pending_commands(self, content: str) -> None:
+        """Optimized command processing"""
+        if not content or not self._pending_commands:
+            return
 
-    def _open_serial_port(self, port, baud_rate):
-        try:
-            self._log(f"Trying to open {port} at {baud_rate} baud.")
-            return serial.Serial(port, baud_rate, timeout=0.05)
-        except serial.SerialException:
-            self._log(f"Failed to open {port} at {baud_rate} baud.")
-            return None
+        with self._command_lock:
+            if not self._pending_commands:
+                return
+                
+            # Get oldest command without min() call
+            oldest_id = next(iter(self._pending_commands))
+            pending = self._pending_commands[oldest_id]
+
+            if pending.future.done():
+                return
+
+            # Fast string comparison
+            if content == pending.command:
+                if not pending.expect_response:
+                    pending.future.set_result(pending.command)
+                    del self._pending_commands[oldest_id]
+            else:
+                pending.future.set_result(content)
+                del self._pending_commands[oldest_id]
+
+    def _cleanup_timed_out_commands(self) -> None:
+        """Optimized cleanup - batch operations"""
+        if not self._pending_commands:
+            return
+            
+        current_time = time.time()
+        with self._command_lock:
+            # Collect timed out commands
+            timed_out = [
+                (cmd_id, pending) 
+                for cmd_id, pending in self._pending_commands.items()
+                if current_time - pending.timestamp > pending.timeout
+            ]
+            
+            # Batch remove
+            for cmd_id, pending in timed_out:
+                del self._pending_commands[cmd_id]
+                if not pending.future.done():
+                    pending.future.set_exception(
+                        MakcuTimeoutError(f"Command #{cmd_id} timed out")
+                    )
+
+    def _listen(self) -> None:
+        """Ultra-optimized listener for gaming performance"""
+        # Pre-allocate buffers
+        read_buffer = bytearray(4096)
+        line_buffer = bytearray(256)
+        line_pos = 0
+        
+        # Cache frequently accessed attributes
+        serial_read = self.serial.read
+        serial_in_waiting = lambda: self.serial.in_waiting
+        is_connected = lambda: self._is_connected
+        stop_requested = self._stop_event.is_set
+        
+        # Timing for cleanup
+        last_cleanup = time.time()
+        cleanup_interval = 0.05  # 50ms cleanup interval
+        
+        while is_connected() and not stop_requested():
+            try:
+                # Check bytes available
+                bytes_available = serial_in_waiting()
+                if not bytes_available:
+                    time.sleep(0.001)  # 1ms sleep to prevent CPU spinning
+                    continue
+                
+                # Read into pre-allocated buffer
+                bytes_read = serial_read(min(bytes_available, 4096))
+                
+                # Process bytes
+                for byte_val in bytes_read:
+                    # Fast button data check
+                    if byte_val < 32 and byte_val not in (0x0D, 0x0A):
+                        self._handle_button_data(byte_val)
+                    else:
+                        # Build line
+                        if byte_val == 0x0A:  # LF
+                            if line_pos > 0:
+                                # Process line without allocation
+                                line = bytes(line_buffer[:line_pos])
+                                line_pos = 0
+                                
+                                if line:
+                                    response = self._parse_response_line(line)
+                                    if response.content:
+                                        self._process_pending_commands(response.content)
+                        elif byte_val != 0x0D:  # Ignore CR
+                            if line_pos < 256:  # Prevent overflow
+                                line_buffer[line_pos] = byte_val
+                                line_pos += 1
+                
+                # Periodic cleanup with reduced frequency
+                current_time = time.time()
+                if current_time - last_cleanup > cleanup_interval:
+                    self._cleanup_timed_out_commands()
+                    last_cleanup = current_time
+                    
+            except serial.SerialException:
+                if self.auto_reconnect:
+                    self._attempt_reconnect()
+                else:
+                    break
+            except Exception:
+                pass  # Silently continue for maximum performance
 
-    def _change_baud_to_4M(self):
+    def _attempt_reconnect(self) -> None:
+        """Fast reconnection attempt"""
+        if self._reconnect_attempts >= self.MAX_RECONNECT_ATTEMPTS:
+            self._is_connected = False
+            return
+        
+        self._reconnect_attempts += 1
+        
+        try:
+            if self.serial and self.serial.is_open:
+                self.serial.close()
+            
+            time.sleep(self.RECONNECT_DELAY)
+            
+            self.port = self.find_com_port()
+            if not self.port:
+                raise MakcuConnectionError("Device not found")
+            
+            # Use write_timeout for faster failure detection
+            self.serial = serial.Serial(
+                self.port, 
+                self.baudrate, 
+                timeout=0.001,  # 1ms read timeout
+                write_timeout=0.01  # 10ms write timeout
+            )
+            self._change_baud_to_4M()
+            
+            if self.send_init:
+                self.serial.write(b"km.buttons(1)\r")
+                self.serial.flush()
+            
+            self._reconnect_attempts = 0
+            
+        except Exception:
+            time.sleep(self.RECONNECT_DELAY)
+
+    def _change_baud_to_4M(self) -> bool:
+        """Optimized baud rate change"""
         if self.serial and self.serial.is_open:
-            self._log("Sending baud rate switch command to 4M.")
-            self.serial.write(self.baud_change_command)
+            self.serial.write(self.BAUD_CHANGE_COMMAND)
             self.serial.flush()
-            time.sleep(0.05)
+            time.sleep(0.02)  # Reduced from 0.05
             self.serial.baudrate = 4000000
             self._current_baud = 4000000
-            self._log("Switched to 4M baud successfully.")
             return True
         return False
 
-
-    def connect(self):
+    def connect(self) -> None:
+        """Optimized connection with minimal overhead"""
         if self._is_connected:
-            self._log("Already connected.")
             return
-        self.serial = self._open_serial_port(self.port, 115200)
-        if not self.serial:
-            raise MakcuConnectionError(f"Failed to connect to {self.port} at 115200.")
-        self._log(f"Connected to {self.port} at 115200.")
-        if not self._change_baud_to_4M():
-            raise MakcuConnectionError("Failed to switch to 4M baud.")
-        self._is_connected = True
-        if self.send_init:
-            with self._lock:
+        
+        if not self.override_port:
+            self.port = self.find_com_port()
+        else:
+            self.port = self._fallback_com_port
+            
+        if not self.port:
+            raise MakcuConnectionError("Makcu device not found")
+        
+        try:
+            # Optimized serial settings for gaming
+            self.serial = serial.Serial(
+                self.port, 
+                115200, 
+                timeout=0.001,  # 1ms timeout
+                write_timeout=0.01,  # 10ms write timeout
+                xonxoff=False,
+                rtscts=False,
+                dsrdtr=False
+            )
+            
+            if not self._change_baud_to_4M():
+                raise MakcuConnectionError("Failed to switch to 4M baud")
+            
+            self._is_connected = True
+            self._reconnect_attempts = 0
+            
+            if self.send_init:
                 self.serial.write(b"km.buttons(1)\r")
                 self.serial.flush()
-                self._log("Sent init command: km.buttons(1)")
-
+            
+            # Start high-priority listener thread
             self._stop_event.clear()
-            self._listener_thread = threading.Thread(target=self._listen, kwargs={"debug": self.debug}, daemon=True)
+            self._listener_thread = threading.Thread(
+                target=self._listen, 
+                daemon=True,
+                name="MakcuListener"
+            )
             self._listener_thread.start()
+            
+        except Exception as e:
+            if self.serial:
+                self.serial.close()
+            raise MakcuConnectionError(f"Failed to connect: {e}")
 
-    def disconnect(self):
+    def disconnect(self) -> None:
+        """Fast disconnect"""
+        self._is_connected = False
+        
         if self.send_init:
             self._stop_event.set()
-            if self._listener_thread:
-                self._listener_thread.join()
-        with self._lock:
-            if self.serial and self.serial.is_open:
-                self.serial.close()
-            self.serial = None
-            self._is_connected = False
-            self._log("Disconnected.")
-
-    def is_connected(self):
-        return self._is_connected
+            if self._listener_thread and self._listener_thread.is_alive():
+                self._listener_thread.join(timeout=0.1)  # Reduced timeout
+        
+        with self._command_lock:
+            for pending in self._pending_commands.values():
+                if not pending.future.done():
+                    pending.future.cancel()
+            self._pending_commands.clear()
+        
+        if self.serial and self.serial.is_open:
+            self.serial.close()
+        self.serial = None
 
-    def send_command(self, command, expect_response=False):
+    def send_command(self, command: str, expect_response: bool = False, 
+                timeout: float = DEFAULT_TIMEOUT) -> Optional[str]:
+        """Optimized command sending for minimal latency"""
         if not self._is_connected or not self.serial or not self.serial.is_open:
-            raise MakcuConnectionError("Serial connection not open.")
-        with self._lock:
-            try:
-                self._pause_listener = True
-                self.serial.reset_input_buffer()
-                self.serial.write(command.encode("ascii") + b"\r\n")
-                self.serial.flush()
-                if expect_response:
-                    response = self.receive_response(sent_command=command)
-                    if not response:
-                        raise MakcuTimeoutError(f"No response from device for command: {command}")
-                    return response
-            finally:
-                self._pause_listener = False
-
+            raise MakcuConnectionError("Not connected")
+        
+        # For commands without response, send and return immediately
+        if not expect_response:
+            self.serial.write(f"{command}\r\n".encode('ascii'))
+            self.serial.flush()
+            return command
+        
+        # Commands with response need tracking
+        cmd_id = self._generate_command_id()
+        tagged_command = f"{command}#{cmd_id}"
+        
+        future = Future()
+        
+        with self._command_lock:
+            self._pending_commands[cmd_id] = PendingCommand(
+                command_id=cmd_id,
+                command=command,
+                future=future,
+                timestamp=time.time(),
+                expect_response=expect_response,
+                timeout=timeout
+            )
+        
+        try:
+            self.serial.write(f"{tagged_command}\r\n".encode('ascii'))
+            self.serial.flush()
+            
+            result = future.result(timeout=timeout)
+            return result.split('#')[0] if '#' in result else result
+            
+        except TimeoutError:
+            raise MakcuTimeoutError(f"Command timed out: {command}")
+        except Exception as e:
+            with self._command_lock:
+                self._pending_commands.pop(cmd_id, None)
+            raise
+
+    async def async_send_command(self, command: str, expect_response: bool = False,
+                               timeout: float = DEFAULT_TIMEOUT) -> Optional[str]:
+        """Async command optimized for gaming"""
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            None, self.send_command, command, expect_response, timeout
+        )
+
+    def is_connected(self) -> bool:
+        """Fast connection check"""
+        return self._is_connected and self.serial is not None and self.serial.is_open
+
+    def set_button_callback(self, callback: Optional[Callable[[MouseButton, bool], None]]) -> None:
+        """Set button callback"""
+        self._button_callback = callback
 
-    def get_button_states(self):
-        return dict(self._button_states)
+    def get_button_states(self) -> Dict[str, bool]:
+        """Get button states with optimized lookup"""
+        return {
+            self.BUTTON_MAP[i]: bool(self._button_states & (1 << i))
+            for i in range(5)
+        }
 
     def get_button_mask(self) -> int:
-        return self._last_mask
-
+        """Direct mask access"""
+        return self._last_button_mask
 
-    def enable_button_monitoring(self, enable: bool = True):
+    def enable_button_monitoring(self, enable: bool = True) -> None:
+        """Fast button monitoring toggle"""
         self.send_command("km.buttons(1)" if enable else "km.buttons(0)")
 
-    def catch_button(self, button: str):
-        command = {
-            "LEFT": "km.catch_ml(0)",
-            "RIGHT": "km.catch_mr(0)",
-            "MIDDLE": "km.catch_mm(0)",
-            "MOUSE4": "km.catch_ms1(0)",
-            "MOUSE5": "km.catch_ms2(0)",
-        }.get(button.upper())
-        if command:
-            self.send_command(command)
-        else:
-            raise ValueError(f"Unsupported button: {button}")
+    # Context managers unchanged but included for completeness
+    async def __aenter__(self):
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, self.connect)
+        return self
 
-    def read_captured_clicks(self, button: str) -> int:
-        command = {
-            "LEFT": "km.catch_ml()",
-            "RIGHT": "km.catch_mr()",
-            "MIDDLE": "km.catch_mm()",
-            "MOUSE4": "km.catch_ms1()",
-            "MOUSE5": "km.catch_ms2()",
-        }.get(button.upper())
-        if command:
-            result = self.send_command(command, expect_response=True)
-            try:
-                return int(result.strip())
-            except Exception:
-                return 0
-        else:
-            raise ValueError(f"Unsupported button: {button}")
-
-    def _listen(self, debug=False):
-        self._log("Started listener thread")
-        button_states = {i: False for i in self.button_map}
-        self._last_mask = 0
-        self._last_callback_time = {bit: 0 for bit in self.button_map}
-
-        while self._is_connected and not self._stop_event.is_set():
-            if self._pause_listener:
-                time.sleep(0.001)
-                continue
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(None, self.disconnect)
 
-            try:
-                byte = self.serial.read(1)
-                if not byte:
-                    continue
-
-                value = byte[0]
-                byte_str = str(byte)
-
-                if not byte_str.startswith("b'\\x"):
-                    continue
-
-                if value != self._last_mask:
-                    if byte_str.startswith("b'\\x00"):
-                        for bit, name in self.button_map.items():
-                            button_states[bit] = False
-                            self._button_states[name] = False
-                            if debug:
-                                print(f"{name} -> False")
-                    else:
-                        for bit, name in self.button_map.items():
-                            is_pressed = bool(value & (1 << bit))
-                            button_states[bit] = is_pressed
-                            self._button_states[name] = is_pressed
-                            if debug:
-                                print(f"{name} -> {is_pressed}")
-
-                    if self._button_callback:
-                        for bit, name in self.button_map.items():
-                            previous = bool(self._last_mask & (1 << bit))
-                            current = bool(value & (1 << bit))
-                            if previous != current:
-                                button_enum = self._button_enum_map.get(bit)
-                                if button_enum:
-                                    self._button_callback(button_enum, current)
-
-                    self._last_mask = value
-
-                    if debug:
-                        pressed = [name for bit, name in self.button_map.items() if button_states[bit]]
-                        button_str = ", ".join(pressed) if pressed else "No buttons pressed"
-                        self._log(f"Byte: {value} (0x{value:02X}) -> {button_str}")
-
-            except serial.SerialException as e:
-                if "ClearCommError failed" not in str(e):
-                    self._log(f"Serial error during listening: {e}")
-                    break
+    def __enter__(self):
+        self.connect()
+        return self
 
-        self._log("Listener thread exiting")
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.disconnect()