Skip to content

Commit 62dc403

Browse files
authored
[simple-websocket] Add stubs for simple-websocket (#14868)
1 parent be34e92 commit 62dc403

File tree

7 files changed

+329
-0
lines changed

7 files changed

+329
-0
lines changed

pyrightconfig.stricter.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@
8888
"stubs/seaborn",
8989
"stubs/setuptools/setuptools",
9090
"stubs/shapely",
91+
"stubs/simple-websocket",
9192
"stubs/tensorflow",
9293
"stubs/tqdm",
9394
"stubs/vobject",
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
version = "1.1.*"
2+
upstream_repository = "https://github.com/miguelgrinberg/simple-websocket"
3+
requires = ["wsproto"]
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from .aiows import AioClient as AioClient, AioServer as AioServer
2+
from .errors import ConnectionClosed as ConnectionClosed, ConnectionError as ConnectionError
3+
from .ws import Client as Client, Server as Server
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import asyncio
2+
import socket
3+
from _typeshed import Incomplete, Unused
4+
from _typeshed.wsgi import WSGIEnvironment
5+
from collections.abc import Awaitable, Callable
6+
from ssl import SSLContext
7+
from typing import Any, Literal, TypedDict, type_check_only
8+
9+
from wsproto import ConnectionType, WSConnection
10+
from wsproto.events import Request
11+
from wsproto.frame_protocol import CloseReason
12+
13+
from .asgi import WebSocketASGI, _SocketDataBase, _SocketDataBytes, _SocketDataProtocol, _SocketDataStr
14+
15+
class AioBase:
16+
subprotocol: str | None
17+
connection_type: ConnectionType
18+
receive_bytes: int
19+
ping_interval: float | None
20+
max_message_size: int | None
21+
pong_received: bool
22+
input_buffer: list[bytes | str]
23+
incoming_message: bytes | str | None
24+
incoming_message_len: int
25+
connected: bool
26+
is_server: bool
27+
close_reason: CloseReason
28+
close_message: str
29+
rsock: asyncio.StreamReader
30+
wsock: asyncio.StreamWriter
31+
event: asyncio.Event
32+
ws: WSConnection | None
33+
task: asyncio.Task[None]
34+
def __init__(
35+
self,
36+
connection_type: ConnectionType | None = None,
37+
receive_bytes: int = 4096,
38+
ping_interval: float | None = None,
39+
max_message_size: int | None = None,
40+
) -> None: ...
41+
async def connect(self) -> None: ...
42+
async def handshake(self) -> None: ...
43+
# data can be antyhing. a special case is made for `bytes`, anything else is converted to `str`.
44+
async def send(self, data: bytes | Any) -> None: ...
45+
async def receive(self, timeout: float | None = None) -> bytes | str | None: ...
46+
async def close(self, reason: CloseReason | None = None, message: str | None = None) -> None: ...
47+
def choose_subprotocol(self, request: Request) -> str | None: ...
48+
49+
@type_check_only
50+
class _AioServerRequest(TypedDict):
51+
# this is `aiohttp.web.Request`
52+
aiohttp: Incomplete
53+
sock: None
54+
headers: None
55+
56+
class AioServer(AioBase):
57+
request: _AioServerRequest
58+
headers: dict[str, Any]
59+
subprotocols: list[str]
60+
is_server: Literal[True]
61+
mode: str
62+
connected: bool
63+
def __init__(
64+
self,
65+
request: _AioServerRequest,
66+
subprotocols: list[str] | None = None,
67+
receive_bytes: int = 4096,
68+
ping_interval: float | None = None,
69+
max_message_size: int | None = None,
70+
) -> None: ...
71+
@classmethod
72+
async def accept(
73+
cls,
74+
# this is `aiohttp.web.Request`
75+
aiohttp=None,
76+
asgi: (
77+
tuple[
78+
WSGIEnvironment,
79+
Callable[[], Awaitable[_SocketDataBytes | _SocketDataStr]],
80+
Callable[[_SocketDataBase | _SocketDataProtocol | _SocketDataBytes | _SocketDataStr], Awaitable[None]],
81+
]
82+
| None
83+
) = None,
84+
sock: socket.socket | None = None,
85+
headers: dict[str, Any] | None = None,
86+
subprotocols: list[str] | None = None,
87+
receive_bytes: int = 4096,
88+
ping_interval: float | None = None,
89+
max_message_size: int | None = None,
90+
) -> WebSocketASGI | AioServer: ...
91+
async def handshake(self) -> None: ...
92+
def choose_subprotocol(self, request: Request) -> str | None: ...
93+
94+
class AioClient(AioBase):
95+
url: str
96+
ssl_context: SSLContext | None
97+
is_secure: bool
98+
host: str
99+
port: int
100+
path: str
101+
subprotocols: list[str]
102+
extra_headeers: list[tuple[bytes, bytes]]
103+
subprotocol: str | None
104+
connected: bool
105+
def __init__(
106+
self,
107+
url: str,
108+
subprotocols: list[str] | None = None,
109+
headers: dict[str, Any] | None = None,
110+
receive_bytes: int = 4096,
111+
ping_interval: float | None = None,
112+
max_message_size: int | None = None,
113+
ssl_context: SSLContext | None = None,
114+
) -> None: ...
115+
# the source code itself has this override
116+
@classmethod
117+
async def connect( # type: ignore[override]
118+
cls,
119+
url: str,
120+
subprotocols: list[str] | None = None,
121+
headers: dict[str, Any] | None = None,
122+
receive_bytes: int = 4096,
123+
ping_interval: float | None = None,
124+
max_message_size: int | None = None,
125+
ssl_context: SSLContext | None = None,
126+
thread_class: Unused = None,
127+
event_class: Unused = None,
128+
) -> AioClient: ...
129+
async def handshake(self) -> None: ...
130+
async def close(self, reason: CloseReason | None = None, message: str | None = None) -> None: ...
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
from _typeshed.wsgi import WSGIEnvironment
2+
from collections.abc import Awaitable, Callable
3+
from typing import TypedDict, type_check_only
4+
5+
@type_check_only
6+
class _SocketDataBase(TypedDict):
7+
type: str
8+
9+
@type_check_only
10+
class _SocketDataProtocol(_SocketDataBase):
11+
subprotocol: str | None
12+
13+
@type_check_only
14+
class _SocketDataStr(_SocketDataBase):
15+
text: str
16+
17+
@type_check_only
18+
class _SocketDataBytes(_SocketDataBase):
19+
bytes: bytes
20+
21+
class WebSocketASGI:
22+
subprotocols: list[str]
23+
subprotocol: str
24+
connected: bool
25+
# this is set in `close` to `False`
26+
conncted: bool
27+
def __init__(
28+
self,
29+
scope: WSGIEnvironment,
30+
receive: Callable[[], Awaitable[_SocketDataBytes | _SocketDataStr]],
31+
send: Callable[[_SocketDataBase | _SocketDataProtocol | _SocketDataBytes | _SocketDataStr], Awaitable[None]],
32+
subprotocols: list[str] | None = None,
33+
) -> None: ...
34+
@classmethod
35+
async def accept(
36+
cls,
37+
scope: WSGIEnvironment,
38+
receive: Callable[[], Awaitable[_SocketDataBytes | _SocketDataStr]],
39+
send: Callable[[_SocketDataBase | _SocketDataProtocol | _SocketDataBytes | _SocketDataStr], Awaitable[None]],
40+
subprotocols: list[str] | None = None,
41+
) -> WebSocketASGI: ...
42+
async def receive(self) -> bytes | str: ...
43+
async def send(self, data: bytes | str) -> None: ...
44+
async def close(self) -> None: ...
Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from wsproto.frame_protocol import CloseReason
2+
3+
class SimpleWebsocketError(RuntimeError): ...
4+
5+
class ConnectionError(SimpleWebsocketError):
6+
status_code: int | None
7+
def __init__(self, status_code: int | None = None) -> None: ...
8+
9+
class ConnectionClosed(SimpleWebsocketError):
10+
reason: CloseReason
11+
message: str | None
12+
def __init__(self, reason: CloseReason = ..., message: str | None = None) -> None: ...
Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
import socket
2+
import threading
3+
from _typeshed import FileDescriptorLike
4+
from _typeshed.wsgi import WSGIEnvironment
5+
from collections.abc import Callable
6+
from selectors import SelectorKey, _EventMask
7+
from ssl import SSLContext
8+
from typing import Any, Protocol, type_check_only
9+
10+
from wsproto import ConnectionType, WSConnection
11+
from wsproto.events import Request
12+
from wsproto.frame_protocol import CloseReason
13+
14+
@type_check_only
15+
class _ThreadClassProtocol(Protocol):
16+
name: str
17+
# this accepts any callable as the target, like `threading.Thread`
18+
def __init__(self, target: Callable[..., Any]) -> None: ...
19+
def start(self) -> None: ...
20+
21+
@type_check_only
22+
class _EventClassProtocol(Protocol):
23+
def clear(self) -> None: ...
24+
def set(self) -> None: ...
25+
def wait(self, timeout: float | None = None) -> bool: ...
26+
27+
@type_check_only
28+
class _SelectorClassProtocol(Protocol):
29+
# the signature of `register` here is the same as `selectors._BaseSelectorImpl` from the stdlib
30+
def register(self, fileobj: FileDescriptorLike, events: _EventMask, data: Any = None) -> SelectorKey: ...
31+
# the signature of `select` here is the same as `selectors.DefaultSelector` from the stdlib
32+
def select(self, timeout: float | None = None) -> list[tuple[SelectorKey, _EventMask]]: ...
33+
def close(self) -> None: ...
34+
35+
class Base:
36+
subprotocol: str | None
37+
sock: socket.socket | None
38+
receive_bytes: int
39+
ping_interval: float | None
40+
max_message_size: int | None
41+
pong_received: bool
42+
input_buffer: list[bytes | str]
43+
incoming_message: bytes | str | None
44+
incoming_message_len: int
45+
connected: bool
46+
is_server: bool
47+
close_reason: CloseReason
48+
close_message: str | None
49+
selector_class: type[_SelectorClassProtocol]
50+
event: _EventClassProtocol | threading.Event
51+
ws: WSConnection
52+
thread: _ThreadClassProtocol | threading.Thread
53+
def __init__(
54+
self,
55+
sock: socket.socket | None = None,
56+
connection_type: ConnectionType | None = None,
57+
receive_bytes: int = 4096,
58+
ping_interval: float | None = None,
59+
max_message_size: int | None = None,
60+
thread_class: type[_ThreadClassProtocol] | None = None,
61+
event_class: type[_EventClassProtocol] | None = None,
62+
selector_class: type[_SelectorClassProtocol] | None = None,
63+
) -> None: ...
64+
def handshake(self) -> None: ...
65+
# data can be antyhing. a special case is made for `bytes`, anything else is converted to `str`.
66+
def send(self, data: bytes | Any) -> None: ...
67+
def receive(self, timeout: float | None = None) -> bytes | str | None: ...
68+
def close(self, reason: CloseReason | None = None, message: str | None = None) -> None: ...
69+
def choose_subprotocol(self, request: Request) -> str | None: ...
70+
71+
class Server(Base):
72+
environ: WSGIEnvironment
73+
subprotocols: list[str]
74+
mode: str
75+
connected: bool
76+
def __init__(
77+
self,
78+
environ: WSGIEnvironment,
79+
subprotocols: list[str] | None = None,
80+
receive_bytes: int = 4096,
81+
ping_interval: float | None = None,
82+
max_message_size: int | None = None,
83+
thread_class: type[_ThreadClassProtocol] | None = None,
84+
event_class: type[_EventClassProtocol] | None = None,
85+
selector_class: type[_SelectorClassProtocol] | None = None,
86+
) -> None: ...
87+
@classmethod
88+
def accept(
89+
cls,
90+
environ: WSGIEnvironment,
91+
subprotocols: list[str] | None = None,
92+
receive_bytes: int = 4096,
93+
ping_interval: float | None = None,
94+
max_message_size: int | None = None,
95+
thread_class: type[_ThreadClassProtocol] | None = None,
96+
event_class: type[_EventClassProtocol] | None = None,
97+
selector_class: type[_SelectorClassProtocol] | None = None,
98+
) -> Server: ...
99+
def handshake(self) -> None: ...
100+
def choose_subprotocol(self, request: Request) -> str | None: ...
101+
102+
class Client(Base):
103+
host: str
104+
port: int
105+
path: str
106+
subprotocols: list[str]
107+
extra_headeers: list[tuple[bytes, bytes]]
108+
subprotocol: str | None
109+
connected: bool
110+
def __init__(
111+
self,
112+
url: str,
113+
subprotocols: list[str] | None = None,
114+
headers: dict[bytes, bytes] | list[tuple[bytes, bytes]] | None = None,
115+
receive_bytes: int = 4096,
116+
ping_interval: float | None = None,
117+
max_message_size: int | None = None,
118+
ssl_context: SSLContext | None = None,
119+
thread_class: type[_ThreadClassProtocol] | None = None,
120+
event_class: type[_EventClassProtocol] | None = None,
121+
) -> None: ...
122+
@classmethod
123+
def connect(
124+
cls,
125+
url: str,
126+
subprotocols: list[str] | None = None,
127+
headers: dict[bytes, bytes] | list[tuple[bytes, bytes]] | None = None,
128+
receive_bytes: int = 4096,
129+
ping_interval: float | None = None,
130+
max_message_size: int | None = None,
131+
ssl_context: SSLContext | None = None,
132+
thread_class: type[_ThreadClassProtocol] | None = None,
133+
event_class: type[_EventClassProtocol] | None = None,
134+
) -> Client: ...
135+
def handshake(self) -> None: ...
136+
def close(self, reason: CloseReason | None = None, message: str | None = None) -> None: ...

0 commit comments

Comments
 (0)