""" This module provides a Pythonic wrapper around the browser's [WebSocket API](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket), enabling two-way communication with WebSocket servers. Use this for real-time applications: - Pythonic interface to browser WebSockets. - Automatic handling of async event handlers. - Support for receiving text (`str`) and binary (`memoryview`) data. - Support for sending text (`str`) and binary (`bytes` and `bytearray`) data. - Compatible with Pyodide and MicroPython. - Works in webworker contexts. - Naming deliberately follows the JavaScript WebSocket API closely for familiarity. See the Python docs for [an explanation of memoryview](https://docs.python.org/3/library/stdtypes.html#memoryview). ```python from pyscript import WebSocket def on_open(event): print("Connected!") ws.send("Hello server") def on_message(event): print(f"Received: {event.data}") def on_close(event): print("Connection closed") ws = WebSocket(url="ws://localhost:8080/") ws.onopen = on_open ws.onmessage = on_message ws.onclose = on_close ``` """ import js from pyscript.ffi import create_proxy from pyscript.util import as_bytearray, is_awaitable def _attach_event_handler(websocket, handler_name, handler_function): """ Given a `websocket`, and `handler_name`, attach the `handler_function` to the `WebSocket` instance, handling both synchronous and asynchronous handler functions. Creates a JavaScript proxy for the handler and wraps async handlers appropriately. Handles the `WebSocketEvent` wrapping for all handlers. """ if is_awaitable(handler_function): async def async_wrapper(event): await handler_function(WebSocketEvent(event)) wrapped_handler = create_proxy(async_wrapper) else: wrapped_handler = create_proxy( lambda event: handler_function(WebSocketEvent(event)) ) # Note: Direct assignment (websocket[handler_name]) fails in Pyodide. setattr(websocket, handler_name, wrapped_handler) class WebSocketEvent: """ A read-only wrapper for [WebSocket event objects](https://developer.mozilla.org/en-US/docs/Web/API/MessageEvent). This class wraps browser WebSocket events and provides convenient access to event properties. It handles the conversion of binary data from JavaScript typed arrays to Python bytes-like objects. The most commonly used property is `event.data`, which contains the message data for "message" events. ```python def on_message(event): # The event is a WebSocketEvent instance. # For text messages. if isinstance(event.data, str): print(f"Text: {event.data}") else: # For binary messages. print(f"Binary: {len(event.data)} bytes") ``` """ def __init__(self, event): """ Create a WebSocketEvent wrapper from an underlying JavaScript `event`. """ self._event = event def __getattr__(self, attr): """ Get an attribute `attr` from the underlying event object. Handles special conversion of binary data from JavaScript typed arrays to Python `memoryview` objects. """ value = getattr(self._event, attr) if attr == "data" and not isinstance(value, str): if hasattr(value, "to_py"): # Pyodide - convert JavaScript typed array to Python. return value.to_py() else: # MicroPython - manually convert JS ArrayBuffer. return memoryview(as_bytearray(value)) return value class WebSocket: """ This class provides a Python-friendly interface to WebSocket connections, handling communication with WebSocket servers. It supports both text and binary data transmission. Access the underlying WebSocket methods and properties directly if needed. However, the wrapper provides a more Pythonic API. If you need to work with the raw JavaScript WebSocket instance, you can access it via the `_js_websocket` attribute. Using textual (`str`) data: ```python from pyscript import WebSocket # Create WebSocket with handlers as arguments. def handle_message(event): print(f"Got: {event.data}") ws = WebSocket( url="ws://echo.websocket.org/", onmessage=handle_message ) # Or assign handlers after creation. def handle_open(event): ws.send("Hello!") ws.onopen = handle_open ``` Using binary (`memoryview`) data: ```python def handle_message(event): if isinstance(event.data, str): print(f"Text: {event.data}") else: # Binary data as memoryview. print(f"Binary: {len(event.data)} bytes") ws = WebSocket(url="ws://example.com/", onmessage=handle_message) # Send binary data. data = bytearray([0x01, 0x02, 0x03]) ws.send(data) ``` Read more about Python's [`memoryview` here](https://docs.python.org/3/library/stdtypes.html#memoryview). """ # WebSocket ready state constants. CONNECTING = 0 OPEN = 1 CLOSING = 2 CLOSED = 3 def __init__(self, url, protocols=None, **handlers): """ Create a new WebSocket connection from the given `url` (`ws://` or `wss://`). Optionally specify `protocols` (a string or a list of protocol strings) and event handlers (`onopen`, `onmessage`, etc.) as keyword arguments. These arguments and naming conventions mirror those of the [underlying JavaScript WebSocket API](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) for familiarity. If you need access to the underlying JavaScript WebSocket instance, you can get it via the `_js_websocket` attribute. ```python # Basic connection. ws = WebSocket(url="ws://localhost:8080/") # With protocol. ws = WebSocket( url="wss://example.com/socket", protocols="chat" ) # With handlers. ws = WebSocket( url="ws://localhost:8080/", onopen=lambda e: print("Connected"), onmessage=lambda e: print(e.data) ) ``` """ # Create underlying JavaScript WebSocket. if protocols: js_websocket = js.WebSocket.new(url, protocols) else: js_websocket = js.WebSocket.new(url) # Set binary type to arraybuffer for easier Python handling. js_websocket.binaryType = "arraybuffer" # Store the underlying WebSocket. # Use object.__setattr__ to bypass our custom __setattr__. object.__setattr__(self, "_js_websocket", js_websocket) # Attach any event handlers passed as keyword arguments. for handler_name, handler in handlers.items(): setattr(self, handler_name, handler) def __getattr__(self, attr): """ Get an attribute `attr` from the underlying WebSocket. This allows transparent access to WebSocket properties like `readyState`, `url`, `bufferedAmount`, etc. """ return getattr(self._js_websocket, attr) def __setattr__(self, attr, value): """ Set an attribute `attr` on the WebSocket to the given `value`. Event handler attributes (`onopen`, `onmessage`, etc.) are specially handled to create proper proxies. Other attributes are set on the underlying WebSocket directly. """ if attr in ["onclose", "onerror", "onmessage", "onopen"]: _attach_event_handler(self._js_websocket, attr, value) else: setattr(self._js_websocket, attr, value) def send(self, data): """ Send `data` through the WebSocket. Accepts both text (`str`) and binary data (`bytes`, `bytearray`, etc.). Binary data is automatically converted to a JavaScript `Uint8Array`. ```python # Send text. ws.send("Hello server!") # Send binary. ws.send(bytes([1, 2, 3, 4])) ws.send(bytearray([5, 6, 7, 8])) ``` !!! warning The WebSocket **must be in the OPEN state to send data**. """ if isinstance(data, str): self._js_websocket.send(data) else: buffer = js.Uint8Array.new(len(data)) for index, byte_value in enumerate(data): buffer[index] = byte_value self._js_websocket.send(buffer) def close(self, code=None, reason=None): """ Close the WebSocket connection. Optionally specify a `code` (`int`) and a `reason` (`str`) for closing the connection. ```python # Normal close. ws.close() # Close with code and reason. ws.close(code=1000, reason="Task completed") ``` Usage and values for `code` and `reasons` [are explained here](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/close). """ if code and reason: self._js_websocket.close(code, reason) elif code: self._js_websocket.close(code) else: self._js_websocket.close()