1""" 2wsproto 3~~~~~~~ 4 5A WebSocket implementation. 6""" 7from typing import Generator, Optional 8 9from .connection import Connection, ConnectionState, ConnectionType 10from .events import Event 11from .handshake import H11Handshake 12from .typing import Headers 13 14__version__ = "1.0.0" 15 16 17class WSConnection: 18 """ 19 Represents the local end of a WebSocket connection to a remote peer. 20 """ 21 22 def __init__(self, connection_type: ConnectionType) -> None: 23 """ 24 Constructor 25 26 :param wsproto.connection.ConnectionType connection_type: Controls 27 whether the library behaves as a client or as a server. 28 """ 29 self.client = connection_type is ConnectionType.CLIENT 30 self.handshake = H11Handshake(connection_type) 31 self.connection: Optional[Connection] = None 32 33 @property 34 def state(self) -> ConnectionState: 35 """ 36 :returns: Connection state 37 :rtype: wsproto.connection.ConnectionState 38 """ 39 if self.connection is None: 40 return self.handshake.state 41 return self.connection.state 42 43 def initiate_upgrade_connection(self, headers: Headers, path: str) -> None: 44 self.handshake.initiate_upgrade_connection(headers, path) 45 46 def send(self, event: Event) -> bytes: 47 """ 48 Generate network data for the specified event. 49 50 When you want to communicate with a WebSocket peer, you should construct 51 an event and pass it to this method. This method will return the bytes 52 that you should send to the peer. 53 54 :param wsproto.events.Event event: The event to generate data for 55 :returns bytes: The data to send to the peer 56 """ 57 data = b"" 58 if self.connection is None: 59 data += self.handshake.send(event) 60 self.connection = self.handshake.connection 61 else: 62 data += self.connection.send(event) 63 return data 64 65 def receive_data(self, data: Optional[bytes]) -> None: 66 """ 67 Feed network data into the connection instance. 68 69 After calling this method, you should call :meth:`events` to see if the 70 received data triggered any new events. 71 72 :param bytes data: Data received from remote peer 73 """ 74 if self.connection is None: 75 self.handshake.receive_data(data) 76 self.connection = self.handshake.connection 77 else: 78 self.connection.receive_data(data) 79 80 def events(self) -> Generator[Event, None, None]: 81 """ 82 A generator that yields pending events. 83 84 Each event is an instance of a subclass of 85 :class:`wsproto.events.Event`. 86 """ 87 yield from self.handshake.events() 88 if self.connection is not None: 89 yield from self.connection.events() 90 91 92__all__ = ("ConnectionType", "WSConnection") 93