from __future__ import annotations import logging import uuid from collections import defaultdict from collections.abc import AsyncIterator from collections.abc import Awaitable from collections.abc import Callable from contextlib import asynccontextmanager from typing import TYPE_CHECKING from typing import Any from pysignalr.exceptions import ServerError from pysignalr.messages import CancelInvocationMessage from pysignalr.messages import CloseMessage from pysignalr.messages import CompletionClientStreamMessage from pysignalr.messages import CompletionMessage from pysignalr.messages import InvocationClientStreamMessage from pysignalr.messages import InvocationMessage from pysignalr.messages import Message from pysignalr.messages import MessageType from pysignalr.messages import PingMessage from pysignalr.messages import StreamInvocationMessage from pysignalr.messages import StreamItemMessage from pysignalr.protocol.json import JSONProtocol from pysignalr.transport.websocket import DEFAULT_CONNECTION_TIMEOUT from pysignalr.transport.websocket import DEFAULT_MAX_SIZE from pysignalr.transport.websocket import DEFAULT_PING_INTERVAL from pysignalr.transport.websocket import DEFAULT_RETRY_COUNT from pysignalr.transport.websocket import DEFAULT_RETRY_MULTIPLIER from pysignalr.transport.websocket import DEFAULT_RETRY_SLEEP from pysignalr.transport.websocket import WebsocketTransport if TYPE_CHECKING: import ssl from pysignalr.protocol.abstract import Protocol from pysignalr.transport.abstract import Transport EmptyCallback = Callable[[], Awaitable[None]] AnyCallback = Callable[[Any], Awaitable[Any | None]] MessageCallback = Callable[[Message], Awaitable[None | Any]] CompletionMessageCallback = Callable[[CompletionMessage], Awaitable[None]] _logger = logging.getLogger('pysignalr.client') class ClientStream: """ Client to server streaming implementation. Refer to https://docs.microsoft.com/en-gb/aspnet/core/signalr/streaming?view=aspnetcore-5.0#client-to-server-streaming for more information. Attributes: transport (Transport): The transport used to send stream items. target (str): The target method name on the server. invocation_id (str): The unique identifier for the stream. """ def __init__(self, transport: Transport, target: str) -> None: self.transport: Transport = transport self.target: str = target self.invocation_id: str = str(uuid.uuid4()) async def send(self, item: Any) -> None: """ Sends the next item to the server. Args: item (Any): The item to send. """ await self.transport.send(StreamItemMessage(self.invocation_id, item)) async def invoke(self) -> None: """ Starts the streaming process. """ await self.transport.send(InvocationClientStreamMessage([self.invocation_id], self.target, [])) async def complete(self) -> None: """ Completes the streaming process. """ await self.transport.send(CompletionClientStreamMessage(self.invocation_id)) class SignalRClient: """ SignalRClient is a client for SignalR that manages connections, sends messages, and handles incoming messages. Attributes: url (str): The URL of the SignalR server. protocol (Protocol): The protocol used for message encoding/decoding. headers (dict[str, str]): Optional HTTP headers to include in the WebSocket handshake. access_token_factory (Callable[[], str] | None): A factory function to provide access tokens. _message_handlers (defaultdict[str, list[MessageCallback | None]]): Handlers for different message types. _stream_handlers (dict[str, tuple[MessageCallback | None, MessageCallback | None, CompletionMessageCallback | None]]): Handlers for stream messages. _invocation_handlers (dict[str, MessageCallback | None]): Handlers for invocation messages. _transport (WebsocketTransport): The transport used for WebSocket communication. _error_callback (CompletionMessageCallback | None): Callback for error messages. """ def __init__( self, url: str, protocol: Protocol | None = None, headers: dict[str, str] | None = None, ping_interval: int = DEFAULT_PING_INTERVAL, connection_timeout: int = DEFAULT_CONNECTION_TIMEOUT, max_size: int | None = DEFAULT_MAX_SIZE, retry_sleep: float = DEFAULT_RETRY_SLEEP, retry_multiplier: float = DEFAULT_RETRY_MULTIPLIER, retry_count: int = DEFAULT_RETRY_COUNT, access_token_factory: Callable[[], str] | None = None, ssl: ssl.SSLContext | None = None, ) -> None: self._url = url self._protocol = protocol or JSONProtocol() self._headers = headers or {} self._access_token_factory = access_token_factory self._ssl = ssl self._message_handlers: defaultdict[str, list[MessageCallback | None]] = defaultdict(list) self._stream_handlers: dict[ str, tuple[MessageCallback | None, MessageCallback | None, CompletionMessageCallback | None] ] = {} self._invocation_handlers: dict[str, MessageCallback | None] = {} self._transport = WebsocketTransport( url=self._url, protocol=self._protocol, callback=self._on_message, headers=self._headers, ping_interval=ping_interval, retry_sleep=retry_sleep, retry_multiplier=retry_multiplier, retry_count=retry_count, connection_timeout=connection_timeout, max_size=max_size, access_token_factory=access_token_factory, ssl=ssl, ) self._error_callback: CompletionMessageCallback | None = None async def run(self) -> None: """ Runs the SignalR client, managing the connection lifecycle. """ await self._transport.run() def on(self, event: str, callback: AnyCallback) -> None: """ Registers a callback function for a specific event. Args: event (str): The event name. callback (AnyCallback): The callback function. """ self._message_handlers[event].append(callback) def on_open(self, callback: EmptyCallback) -> None: """ Registers a callback function to be called when the connection is opened. Args: callback (EmptyCallback): The callback function. """ self._transport.on_open(callback) def on_close(self, callback: EmptyCallback) -> None: """ Registers a callback function to be called when the connection is closed. Args: callback (EmptyCallback): The callback function. """ self._transport.on_close(callback) def on_error(self, callback: CompletionMessageCallback) -> None: """ Registers a callback function to be called when an error occurs. Args: callback (CompletionMessageCallback): The callback function. """ self._error_callback = callback async def send( self, method: str, arguments: list[dict[str, Any]], on_invocation: MessageCallback | None = None, ) -> None: """ Sends a message to the server. Args: method (str): The method name to invoke on the server. arguments (list[dict[str, Any]]): The arguments to pass to the method. on_invocation (MessageCallback | None): Optional callback for the invocation response. """ invocation_id = str(uuid.uuid4()) message = InvocationMessage(invocation_id, method, arguments, self._headers) self._invocation_handlers[invocation_id] = on_invocation await self._transport.send(message) async def stream( self, event: str, event_params: list[str], on_next: MessageCallback | None = None, on_complete: MessageCallback | None = None, on_error: CompletionMessageCallback | None = None, ) -> None: """ Starts a streaming invocation. Args: event (str): The event name to stream. event_params (list[str]): The parameters for the event. on_next (MessageCallback | None): Optional callback for each stream item. on_complete (MessageCallback | None): Optional callback when the stream is completed. on_error (CompletionMessageCallback | None): Optional callback for errors. """ invocation_id = str(uuid.uuid4()) message = StreamInvocationMessage(invocation_id, event, event_params, self._headers) self._stream_handlers[invocation_id] = (on_next, on_complete, on_error) await self._transport.send(message) @asynccontextmanager async def client_stream(self, target: str) -> AsyncIterator[ClientStream]: """ Context manager for client-to-server streaming. Args: target (str): The target method name on the server. Yields: ClientStream: The client stream instance. """ stream = ClientStream(self._transport, target) await stream.invoke() yield stream await stream.complete() async def _on_message(self, message: Message) -> None: """ Handles incoming messages and routes them to the appropriate handlers. Args: message (Message): The incoming message. """ if message.type == MessageType.invocation_binding_failure: # type: ignore[attr-defined] raise ServerError(str(message)) elif isinstance(message, PingMessage): pass elif isinstance(message, InvocationMessage): await self._on_invocation_message(message) elif isinstance(message, CloseMessage): await self._on_close_message(message) elif isinstance(message, CompletionMessage): await self._on_completion_message(message) elif isinstance(message, StreamItemMessage): await self._on_stream_item_message(message) elif isinstance(message, StreamInvocationMessage): pass elif isinstance(message, CancelInvocationMessage): await self._on_cancel_invocation_message(message) else: raise NotImplementedError async def _on_invocation_message(self, message: InvocationMessage) -> None: """ Handles invocation messages. Args: message (InvocationMessage): The invocation message. """ expects_response = message.invocation_id is not None callbacks = [callback for callback in self._message_handlers[message.target] if callback] if not callbacks: # There are no callbacks for the message.target _logger.warning("No client method with the name '%s' found.", message.target) if expects_response: _logger.error( "No result given for '%s' method and invocation ID '%s'.", message.target, message.invocation_id, ) await self._transport.send( CompletionMessage(invocation_id=message.invocation_id, error="Client didn't provide a result.") ) return None if expects_response and len(callbacks) > 1: # There are multiple callbacks, so multiple results for the message.target _logger.error("Multiple results provided for '%s'. Sending error to server.", message.target) await self._transport.send( CompletionMessage(invocation_id=message.invocation_id, error='Client provided multiple results.') ) return None for callback in callbacks: try: res = await callback(message.arguments) if res: if expects_response: await self._transport.send(CompletionMessage(invocation_id=message.invocation_id, result=res)) else: _logger.warning( "Result given for '%s' method but server is not expecting a result.", message.target ) elif expects_response: _logger.error( "No result given for '%s' method and invocation ID '%s'.", message.target, message.invocation_id, ) await self._transport.send( CompletionMessage( invocation_id=message.invocation_id, error="Client didn't provide a result.", ) ) except Exception as exc: _logger.error("A callback for the method '%s' threw error '%s'.", message.target, exc) if not expects_response: raise exc await self._transport.send(CompletionMessage(invocation_id=message.invocation_id, error=str(exc))) return None async def _on_completion_message(self, message: CompletionMessage) -> None: """ Handles completion messages. Args: message (CompletionMessage): The completion message. """ if message.error: if self._error_callback is None: raise RuntimeError('Error callback is not set') await self._error_callback(message) callback = self._invocation_handlers.pop(message.invocation_id) if callback is not None: await callback(message) async def _on_stream_item_message(self, message: StreamItemMessage) -> None: """ Handles stream item messages. Args: message (StreamItemMessage): The stream item message. """ callback, _, _ = self._stream_handlers[message.invocation_id] if callback: await callback(message.item) async def _on_cancel_invocation_message(self, message: CancelInvocationMessage) -> None: """ Handles cancel invocation messages. Args: message (CancelInvocationMessage): The cancel invocation message. """ _, _, callback = self._stream_handlers[message.invocation_id] if callback: await callback(message) # type: ignore[arg-type] async def _on_close_message(self, message: CloseMessage) -> None: """ Handles close messages. Args: message (CloseMessage): The close message. """ if message.error: raise ServerError(message.error)