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