Repository: encode/starlette

Issue:
[Feature request] starlette.requests.Request should work for websocket scopes
Following use-case:
I am writing some authentication middleware for websockets. I'm using session tokens that are stored as cookie. For that, the relevant data in scope object needs to be rehashed. This functionality already exist within Starlette in the form of starlette.requests.Request, but it only works for http requests. Perhaps there is some rationale behind this, but it eludes me for now.


Discussion:
Cookie support ought to be on WebSocket too - it’s just an oversight. Really we ought to include a common base class, HTTPConnection.

Resolving patch:
diff --git a/starlette/requests.py b/starlette/requests.py
--- a/starlette/requests.py
+++ b/starlette/requests.py
@@ -17,16 +17,15 @@ class ClientDisconnect(Exception):
     pass
 
 
-async def empty_receive() -> Message:
-    raise RuntimeError("Receive channel has not been made available")
+class HTTPConnection(Mapping):
+    """
+    A base class for incoming HTTP connections, that is used to provide
+    any functionality that is common to both `Request` and `WebSocket`.
+    """
 
-
-class Request(Mapping):
     def __init__(self, scope: Scope, receive: Receive = None) -> None:
-        assert scope["type"] == "http"
+        assert scope["type"] in ("http", "websocket")
         self._scope = scope
-        self._receive = empty_receive if receive is None else receive
-        self._stream_consumed = False
 
     def __getitem__(self, key: str) -> str:
         return self._scope[key]
@@ -37,10 +36,6 @@ def __iter__(self) -> typing.Iterator[str]:
     def __len__(self) -> int:
         return len(self._scope)
 
-    @property
-    def method(self) -> str:
-        return self._scope["method"]
-
     @property
     def url(self) -> URL:
         if not hasattr(self, "_url"):
@@ -104,15 +99,31 @@ def user(self) -> typing.Any:
         ), "AuthenticationMiddleware must be installed to access request.user"
         return self._scope["user"]
 
-    @property
-    def receive(self) -> Receive:
-        return self._receive
-
     def url_for(self, name: str, **path_params: typing.Any) -> str:
         router = self._scope["router"]
         url_path = router.url_path_for(name, **path_params)
         return url_path.make_absolute_url(base_url=self.url)
 
+
+async def empty_receive() -> Message:
+    raise RuntimeError("Receive channel has not been made available")
+
+
+class Request(HTTPConnection):
+    def __init__(self, scope: Scope, receive: Receive = empty_receive):
+        super().__init__(scope)
+        assert scope["type"] == "http"
+        self._receive = receive
+        self._stream_consumed = False
+
+    @property
+    def method(self) -> str:
+        return self._scope["method"]
+
+    @property
+    def receive(self) -> Receive:
+        return self._receive
+
     async def stream(self) -> typing.AsyncGenerator[bytes, None]:
         if hasattr(self, "_body"):
             yield self._body
diff --git a/starlette/websockets.py b/starlette/websockets.py
--- a/starlette/websockets.py
+++ b/starlette/websockets.py
@@ -4,6 +4,7 @@
 from collections.abc import Mapping
 
 from starlette.datastructures import URL, Headers, QueryParams
+from starlette.requests import HTTPConnection
 from starlette.types import Message, Receive, Scope, Send
 
 
@@ -18,50 +19,50 @@ def __init__(self, code: int = 1000) -> None:
         self.code = code
 
 
-class WebSocket(Mapping):
+class WebSocket(HTTPConnection):
     def __init__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        super().__init__(scope)
         assert scope["type"] == "websocket"
-        self._scope = scope
         self._receive = receive
         self._send = send
         self.client_state = WebSocketState.CONNECTING
         self.application_state = WebSocketState.CONNECTING
 
-    def __getitem__(self, key: str) -> str:
-        return self._scope[key]
-
-    def __iter__(self) -> typing.Iterator:
-        return iter(self._scope)
-
-    def __len__(self) -> int:
-        return len(self._scope)
-
-    @property
-    def url(self) -> URL:
-        if not hasattr(self, "_url"):
-            self._url = URL(scope=self._scope)
-        return self._url
-
-    @property
-    def headers(self) -> Headers:
-        if not hasattr(self, "_headers"):
-            self._headers = Headers(scope=self._scope)
-        return self._headers
-
-    @property
-    def query_params(self) -> QueryParams:
-        if not hasattr(self, "_query_params"):
-            self._query_params = QueryParams(scope=self._scope)
-        return self._query_params
-
-    @property
-    def path_params(self) -> dict:
-        return self._scope.get("path_params", {})
-
-    def url_for(self, name: str, **path_params: typing.Any) -> str:
-        router = self._scope["router"]
-        url_path = router.url_path_for(name, **path_params)
-        return url_path.make_absolute_url(base_url=self.url)
+    # def __getitem__(self, key: str) -> str:
+    #     return self._scope[key]
+    #
+    # def __iter__(self) -> typing.Iterator:
+    #     return iter(self._scope)
+    #
+    # def __len__(self) -> int:
+    #     return len(self._scope)
+    #
+    # @property
+    # def url(self) -> URL:
+    #     if not hasattr(self, "_url"):
+    #         self._url = URL(scope=self._scope)
+    #     return self._url
+    #
+    # @property
+    # def headers(self) -> Headers:
+    #     if not hasattr(self, "_headers"):
+    #         self._headers = Headers(scope=self._scope)
+    #     return self._headers
+    #
+    # @property
+    # def query_params(self) -> QueryParams:
+    #     if not hasattr(self, "_query_params"):
+    #         self._query_params = QueryParams(scope=self._scope)
+    #     return self._query_params
+    #
+    # @property
+    # def path_params(self) -> dict:
+    #     return self._scope.get("path_params", {})
+    #
+    # def url_for(self, name: str, **path_params: typing.Any) -> str:
+    #     router = self._scope["router"]
+    #     url_path = router.url_path_for(name, **path_params)
+    #     return url_path.make_absolute_url(base_url=self.url)
 
     async def receive(self) -> Message:
         """


Test patch: