diff --git a/responder/models.py b/responder/models.py index ebce351..b3a4c56 100644 --- a/responder/models.py +++ b/responder/models.py @@ -13,6 +13,7 @@ from requests.cookies import RequestsCookieJar from starlette.datastructures import MutableHeaders from starlette.requests import Request as StarletteRequest from starlette.responses import Response as StarletteResponse +from starlette.websockets import WebSocket as StarletteWebSocket from urllib.parse import parse_qs @@ -287,3 +288,54 @@ class Response: body, status_code=self.status_code, headers=headers ) await response(receive, send) + + +class WebSocket: + # TODO: __slots__ + def __init__(self, scope, receive, send): + self._starlette = StarletteWebSocket(scope, receive, send) + + headers = CaseInsensitiveDict() + for header, value in self._starlette.headers.items(): + headers[header] = value + + self._headers = headers + + @property + def url(self): + return rfc3986.urlparse(str(self._starlette.url)) + + @property + def params(self): + """A dictionary of the parsed query parameters used for the Request.""" + try: + return QueryDict(self.url.query) + except AttributeError: + return QueryDict({}) + + async def accept(self): + return (await self._starlette.accept()) + + async def receive(self): + return (await self._starlette.receive()) + + async def content(self): + return (await self._starlette.receive_bytes()) + + async def text(self): + return (await self._starlette.receive_text()) + + async def send(self): + await self._starlette.send() + + async def send_text(self, data): + await self._starlette.send_text(data) + + async def send_json(self, data): + await self._starlette.send_json(data) + + async def send_media(self, data): + await self.send_json(data) + + async def close(self, code=1000) -> None: + await self._starlette.close(code)