jupyverse 0.10.1__tar.gz → 0.10.2__tar.gz

{jupyverse-0.10.1 → jupyverse-0.10.2}/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Version history
 
+## 0.10.2
+
+- Rework `test_kernels.py` to use ASGI transport.
+
 ## 0.10.1
 
 - Fix optional dependencies pinning.

{jupyverse-0.10.1 → jupyverse-0.10.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jupyverse
-Version: 0.10.1
+Version: 0.10.2
 Summary: A set of FPS plugins implementing a Jupyter server
 Project-URL: Homepage, https://jupyter.org
 Author-email: Jupyter Development Team <jupyter@googlegroups.com>
@@ -13,7 +13,7 @@ Requires-Dist: fps-contents<0.11.0,>=0.10.0
 Requires-Dist: fps-file-id<0.3.0,>=0.2.0
 Requires-Dist: fps-frontend<0.10.0,>=0.9.0
 Requires-Dist: fps-kernel-subprocess<0.2.0,>=0.1.0
-Requires-Dist: fps-kernels<0.10.0,>=0.9.0
+Requires-Dist: fps-kernels<0.10.0,>=0.9.1
 Requires-Dist: fps-lab<0.10.0,>=0.9.0
 Requires-Dist: fps-nbconvert<0.10.0,>=0.9.0
 Requires-Dist: fps-terminals<0.10.0,>=0.9.0

{jupyverse-0.10.1 → jupyverse-0.10.2}/plugins/kernels/fps_kernels/kernel_server/server.py

@@ -6,7 +6,7 @@ from datetime import datetime, timezone
 
 from anyio import TASK_STATUS_IGNORED, Event, create_task_group, move_on_after
 from anyio.abc import TaskStatus
-from fastapi import WebSocket, WebSocketDisconnect
+from fastapi import WebSocket
 from starlette.websockets import WebSocketState
 
 from jupyverse_api.kernel import DefaultKernelFactory, KernelFactory
@@ -160,7 +160,7 @@ class KernelServer:
     async def listen_web(self, websocket: AcceptedWebSocket, stop_event: Event):
         try:
             await self.send_to_kernel(websocket)
-        except WebSocketDisconnect:
+        except BaseException:
             pass
         finally:
             stop_event.set()

{jupyverse-0.10.1 → jupyverse-0.10.2}/plugins/kernels/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "fps_kernels"
-version = "0.9.0"
+version = "0.9.1"
 description = "An FPS plugin for the kernels API"
 keywords = ["jupyter", "server", "fastapi", "plugins"]
 requires-python = ">=3.9"

{jupyverse-0.10.1 → jupyverse-0.10.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "jupyverse"
-version = "0.10.1"
+version = "0.10.2"
 description = "A set of FPS plugins implementing a Jupyter server"
 keywords = ["jupyter", "server", "fastapi", "plugins"]
 requires-python = ">=3.9"
@@ -13,7 +13,7 @@ dependencies = [
     "fps-contents >=0.10.0,<0.11.0",
     "fps-file-id >=0.2.0,<0.3.0",
     "fps-kernel-subprocess >=0.1.0,<0.2.0",
-    "fps-kernels >=0.9.0,<0.10.0",
+    "fps-kernels >=0.9.1,<0.10.0",
     "fps-terminals >=0.9.0,<0.10.0",
     "fps-nbconvert >=0.9.0,<0.10.0",
     "fps-yjs >=0.10.0,<0.11.0",

{jupyverse-0.10.1 → jupyverse-0.10.2}/tests/test_kernels.py

@@ -8,6 +8,7 @@ from fps import get_root_module, merge_config
 from fps_kernels.kernel_server.server import KernelServer, kernels
 from httpx import AsyncClient
 from httpx_ws import aconnect_ws
+from utils import ASGIWebSocketTransport
 
 from jupyverse_api.kernel import Kernel
 
@@ -16,6 +17,9 @@ os.environ["PYDEVD_DISABLE_FILE_VALIDATION"] = "1"
 CONFIG = {
     "jupyverse": {
         "type": "jupyverse",
+            "config": {
+                "start_server": False,
+            },
         "modules": {
             "app": {
                 "type": "app",
@@ -29,9 +33,6 @@ CONFIG = {
             "contents": {
                 "type": "contents",
             },
-            "file_id": {
-                "type": "file_id",
-            },
             "frontend": {
                 "type": "frontend",
             },
@@ -47,9 +48,6 @@ CONFIG = {
             "kernels": {
                 "type": "kernels",
             },
-            "yjs": {
-                "type": "yjs",
-            },
         },
     }
 }
@@ -57,87 +55,96 @@ CONFIG = {
 
 @pytest.mark.anyio
 @pytest.mark.parametrize("auth_mode", ("noauth",))
-async def test_kernel_messages(auth_mode, capfd, unused_tcp_port):
+async def test_kernel_messages(auth_mode, capfd):
     kernel_id = "kernel_id_0"
     kernel_name = "python3"
     kernelspec_path = (
         Path(sys.prefix) / "share" / "jupyter" / "kernels" / kernel_name / "kernel.json"
     )
     assert kernelspec_path.exists()
-    async with create_task_group() as tg:
-        msg = {
-            "channel": "shell",
-            "parent_header": None,
-            "content": None,
-            "metadata": None,
-            "header": {
-                "msg_type": "msg_type_0",
-                "msg_id": "0",
-            },
-        }
-
-        config = merge_config(
-            CONFIG,
-            {
-                "jupyverse": {
-                    "config": {"port": unused_tcp_port},
-                    "modules": {
-                        "auth": {
-                            "config": {
-                                "mode": auth_mode,
-                            }
+    msg = {
+        "channel": "shell",
+        "parent_header": None,
+        "content": None,
+        "metadata": None,
+        "header": {
+            "msg_type": "msg_type_0",
+            "msg_id": "0",
+        },
+    }
+
+    config = merge_config(
+        CONFIG,
+        {
+            "jupyverse": {
+                "modules": {
+                    "auth": {
+                        "config": {
+                            "mode": auth_mode,
                         }
-                    },
-                }
-            },
-        )
-        async with get_root_module(config), AsyncClient():
+                    }
+                },
+            }
+        },
+    )
+    async with get_root_module(config) as root_module:
+        app = root_module.app
+        transport = ASGIWebSocketTransport(app=app)
+        async with AsyncClient(transport=transport, base_url="http://testserver") as client:
             kernel_server = KernelServer(
                 kernelspec_path=kernelspec_path,
                 capture_kernel_output=False,
                 default_kernel_factory=KernelFactory(),
             )
-            await tg.start(kernel_server.start)
-            kernels[kernel_id] = {"server": kernel_server, "driver": None}
-            # block msg_type_0
-            kernel_server.block_messages("msg_type_0")
-            async with aconnect_ws(
-                f"http://127.0.0.1:{unused_tcp_port}/api/kernels/{kernel_id}/channels?session_id=session_id_0",
-            ) as websocket:
-                await websocket.send_json(msg)
-            await sleep(0.1)
-            out, err = capfd.readouterr()
-            assert not err
-
-            # allow only msg_type_0
-            kernel_server.allow_messages("msg_type_0")
-            async with aconnect_ws(
-                f"http://127.0.0.1:{unused_tcp_port}/api/kernels/{kernel_id}/channels?session_id=session_id_0",
-            ) as websocket:
-                await websocket.send_json(msg)
-            await sleep(0.1)
-            out, err = capfd.readouterr()
-            assert err.count("msg_type_0") == 1
-
-            # block all messages
-            kernel_server.allow_messages([])
-            async with aconnect_ws(
-                f"http://127.0.0.1:{unused_tcp_port}/api/kernels/{kernel_id}/channels?session_id=session_id_0",
-            ) as websocket:
-                await websocket.send_json(msg)
-            await sleep(0.1)
-            out, err = capfd.readouterr()
-            assert not err
-
-            # allow all messages
-            kernel_server.allow_messages()
-            async with aconnect_ws(
-                f"http://127.0.0.1:{unused_tcp_port}/api/kernels/{kernel_id}/channels?session_id=session_id_0",
-            ) as websocket:
-                await websocket.send_json(msg)
-            await sleep(0.1)
-            out, err = capfd.readouterr()
-            assert err.count("msg_type_0") >= 1
+            async with create_task_group() as tg:
+                await tg.start(kernel_server.start)
+                kernels[kernel_id] = {"server": kernel_server, "driver": None}
+
+                # block msg_type_0
+                kernel_server.block_messages("msg_type_0")
+                async with aconnect_ws(
+                    f"http://testserver/api/kernels/{kernel_id}/channels?session_id=session_id_0",
+                    client,
+                ) as websocket:
+                    await websocket.send_json(msg)
+                await sleep(0.1)
+                out, err = capfd.readouterr()
+                assert not err
+
+                # allow only msg_type_0
+                kernel_server.allow_messages("msg_type_0")
+                async with aconnect_ws(
+                    f"http://testserver/api/kernels/{kernel_id}/channels?session_id=session_id_0",
+                    client,
+                ) as websocket:
+                    await websocket.send_json(msg)
+                    await sleep(0.1)
+                out, err = capfd.readouterr()
+                assert err.count("msg_type_0") == 1
+
+                # block all messages
+                kernel_server.allow_messages([])
+                async with aconnect_ws(
+                    f"http://testserver/api/kernels/{kernel_id}/channels?session_id=session_id_0",
+                    client,
+                ) as websocket:
+                    await websocket.send_json(msg)
+                await sleep(0.1)
+                out, err = capfd.readouterr()
+                assert not err
+
+                # allow all messages
+                kernel_server.allow_messages()
+                async with aconnect_ws(
+                    f"http://testserver/api/kernels/{kernel_id}/channels?session_id=session_id_0",
+                    client,
+                ) as websocket:
+                    await websocket.send_json(msg)
+                await sleep(0.1)
+                out, err = capfd.readouterr()
+                assert err.count("msg_type_0") >= 1
+
+                await kernel_server.stop()
 
 
 class KernelFactory:

jupyverse-0.10.2/tests/utils.py

@@ -0,0 +1,385 @@
+import contextlib
+import queue
+import typing
+from types import TracebackType
+from typing import Optional
+from uuid import uuid4
+
+import anyio
+import wsproto
+from anyio import Lock
+from httpcore import AsyncNetworkStream
+from httpx import ASGITransport, AsyncByteStream, Request, Response
+from httpx_ws._exceptions import WebSocketDisconnect
+from wsproto.frame_protocol import CloseReason
+
+
+async def authenticate_client(http, port, permissions={}):
+    # create a new user
+    username = uuid4().hex
+    # if logged in, log out
+    first_time = True
+    while True:
+        response = await http.get(f"http://127.0.0.1:{port}/api/me")
+        if response.status_code == 403:
+            break
+        assert first_time
+        response = await http.post(f"http://127.0.0.1:{port}/auth/logout")
+        assert response.status_code == 200
+        first_time = False
+
+    # register user
+    register_body = {
+        "email": f"{username}@example.com",
+        "password": username,
+        "username": username,
+        "permissions": permissions,
+    }
+    response = await http.post(f"http://127.0.0.1:{port}/auth/register", json=register_body)
+    # check that we cannot register if not logged in
+    assert response.status_code == 403
+    # log in as admin
+    login_body = {"username": "admin@jupyter.com", "password": "jupyverse"}
+    response = await http.post(f"http://127.0.0.1:{port}/auth/login", data=login_body)
+    assert response.status_code == 204
+    # register user
+    response = await http.post(f"http://127.0.0.1:{port}/auth/register", json=register_body)
+    assert response.status_code == 201
+
+    # log out
+    response = await http.post(f"http://127.0.0.1:{port}/auth/logout")
+    assert response.status_code == 204
+    # check that we can't get our identity, since we're not logged in
+    response = await http.get(f"http://127.0.0.1:{port}/api/me")
+    assert response.status_code == 403
+
+    # log in with registered user
+    login_body = {"username": f"{username}@example.com", "password": username}
+    response = await http.post(f"http://127.0.0.1:{port}/auth/login", data=login_body)
+    assert response.status_code == 204
+    # we should now have a cookie
+    assert "fastapiusersauth" in http.cookies
+    # check our identity, since we're logged in
+    response = await http.get(
+        f"http://127.0.0.1:{port}/api/me", params={"permissions": permissions}
+    )
+    assert response.status_code == 200
+    me = response.json()
+    assert me["identity"]["username"] == username
+    # check our permissions
+    assert me["permissions"] == permissions
+
+
+def create_content(
+    content: Optional[list],
+    type: str,
+    size: Optional[int],
+    mimetype: Optional[str],
+    name: str,
+    path: str,
+    format: Optional[str],
+) -> dict:
+    return {
+        "content": content,
+        "created": None,
+        "format": format,
+        "last_modified": None,
+        "mimetype": mimetype,
+        "name": name,
+        "path": path,
+        "size": size,
+        "type": type,
+        "writable": True,
+    }
+
+
+def clear_content_values(content: dict, keys: list[str] = []):
+    for k in content:
+        if k in keys:
+            content[k] = None
+        if k == "content" and isinstance(content[k], list):
+            for c in content[k]:
+                clear_content_values(c, keys)
+    return content
+
+
+def sort_content_by_name(content: dict):
+    for k in content:
+        if k == "content" and isinstance(content[k], list):
+            # FIXME: this sorting algorithm is terrible!
+            names = [c["name"] for c in content[k]]
+            names.sort()
+            new_content = []
+            for name in names:
+                for i, c in enumerate(content[k]):
+                    if c["name"] == name:
+                        break
+                content[k].pop(i)
+                new_content.append(c)
+            content[k] = new_content
+            for c in content[k]:
+                sort_content_by_name(c)
+    return content
+
+
+class Websocket:
+    def __init__(self, websocket, path: str):
+        self._websocket = websocket
+        self._path = path
+        self._send_lock = Lock()
+
+    @property
+    def path(self) -> str:
+        return self._path
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> bytes:
+        try:
+            message = await self.recv()
+        except Exception:
+            raise StopAsyncIteration()
+        return message
+
+    async def send(self, message: bytes):
+        async with self._send_lock:
+            await self._websocket.send_bytes(message)
+
+    async def recv(self) -> bytes:
+        b = await self._websocket.receive_bytes()
+        return bytes(b)
+
+
+Scope = dict[str, typing.Any]
+Message = dict[str, typing.Any]
+Receive = typing.Callable[[], typing.Awaitable[Message]]
+Send = typing.Callable[[Scope], typing.Coroutine[None, None, None]]
+ASGIApp = typing.Callable[[Scope, Receive, Send], typing.Coroutine[None, None, None]]
+
+
+class ASGIWebSocketTransportError(Exception):
+    pass
+
+
+class UnhandledASGIMessageType(ASGIWebSocketTransportError):
+    def __init__(self, message: Message) -> None:
+        self.message = message
+
+
+class UnhandledWebSocketEvent(ASGIWebSocketTransportError):
+    def __init__(self, event: wsproto.events.Event) -> None:
+        self.event = event
+
+
+class ASGIWebSocketAsyncNetworkStream(AsyncNetworkStream):
+    def __init__(
+        self, app: ASGIApp, scope: Scope, task_group: anyio.abc.TaskGroup
+    ) -> None:
+        self.app = app
+        self.scope = scope
+        self._task_group = task_group
+        self._receive_queue: queue.Queue[Message] = queue.Queue()
+        self._send_queue: queue.Queue[Message] = queue.Queue()
+        self.connection = wsproto.WSConnection(wsproto.ConnectionType.SERVER)
+        self.connection.initiate_upgrade_connection(scope["headers"], scope["path"])
+        self._aentered = False
+
+    async def __aenter__(
+        self,
+    ) -> tuple["ASGIWebSocketAsyncNetworkStream", bytes]:
+        if self._aentered:
+            raise RuntimeError(
+                "Cannot use ASGIWebSocketAsyncNetworkStream in a context manager twice"
+            )
+        self._aentered = True
+        self._task_group.start_soon(self._run)
+        async with contextlib.AsyncExitStack() as stack:
+            await self.send({"type": "websocket.connect"})
+            message = await self.receive()
+
+            stack.push_async_callback(self.aclose)
+
+            if message["type"] == "websocket.close":
+                await stack.pop_all().aclose()
+                raise WebSocketDisconnect(message["code"], message.get("reason"))
+
+            assert message["type"] == "websocket.accept"
+            retval = self, self._build_accept_response(message)
+            self._exit_stack = stack.pop_all()
+        return retval
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> typing.Union[bool, None]:
+        return await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def read(
+        self, max_bytes: int, timeout: typing.Optional[float] = None
+    ) -> bytes:
+        message: Message = await self.receive(timeout=timeout)
+        type = message["type"]
+
+        if type not in {"websocket.send", "websocket.close"}:
+            raise UnhandledASGIMessageType(message)
+
+        event: wsproto.events.Event
+        if type == "websocket.send":
+            data_str: typing.Optional[str] = message.get("text")
+            if data_str is not None:
+                event = wsproto.events.TextMessage(data_str)
+            data_bytes: typing.Optional[bytes] = message.get("bytes")
+            if data_bytes is not None:
+                event = wsproto.events.BytesMessage(data_bytes)
+        elif type == "websocket.close":
+            event = wsproto.events.CloseConnection(message["code"], message["reason"])
+
+        return self.connection.send(event)
+
+    async def write(
+        self, buffer: bytes, timeout: typing.Optional[float] = None
+    ) -> None:
+        self.connection.receive_data(buffer)
+        for event in self.connection.events():
+            if isinstance(event, wsproto.events.Request):
+                pass
+            elif isinstance(event, wsproto.events.CloseConnection):
+                await self.send(
+                    {
+                        "type": "websocket.close",
+                        "code": event.code,
+                        "reason": event.reason,
+                    }
+                )
+            elif isinstance(event, wsproto.events.TextMessage):
+                await self.send({"type": "websocket.receive", "text": event.data})
+            elif isinstance(event, wsproto.events.BytesMessage):
+                await self.send({"type": "websocket.receive", "bytes": event.data})
+            else:
+                raise UnhandledWebSocketEvent(event)
+
+    async def aclose(self) -> None:
+        await self.send({"type": "websocket.close"})
+
+    async def send(self, message: Message) -> None:
+        self._receive_queue.put(message)
+
+    async def receive(self, timeout: typing.Optional[float] = None) -> Message:
+        while self._send_queue.empty():
+            await anyio.sleep(0)
+        return self._send_queue.get(timeout=timeout)
+
+    async def _run(self) -> None:
+        """
+        The sub-thread in which the websocket session runs.
+        """
+        scope = self.scope
+        receive = self._asgi_receive
+        send = self._asgi_send
+        try:
+            await self.app(scope, receive, send)
+        except Exception as e:
+            message = {
+                "type": "websocket.close",
+                "code": CloseReason.INTERNAL_ERROR,
+                "reason": str(e),
+            }
+            await self._asgi_send(message)
+
+    async def _asgi_receive(self) -> Message:
+        while self._receive_queue.empty():
+            await anyio.sleep(0)
+        return self._receive_queue.get()
+
+    async def _asgi_send(self, message: Message) -> None:
+        self._send_queue.put(message)
+
+    def _build_accept_response(self, message: Message) -> bytes:
+        subprotocol = message.get("subprotocol", None)
+        headers = message.get("headers", [])
+        return self.connection.send(
+            wsproto.events.AcceptConnection(
+                subprotocol=subprotocol,
+                extra_headers=headers,
+            )
+        )
+
+
+class ASGIWebSocketTransport(ASGITransport):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.exit_stacks: list[contextlib.AsyncExitStack] = []
+
+    async def __aenter__(self) -> "ASGIWebSocketTransport":
+        async with contextlib.AsyncExitStack() as stack:
+            self._task_group = await stack.enter_async_context(
+                anyio.create_task_group()
+            )
+            self.exit_stack = stack.pop_all()
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: typing.Optional[type[BaseException]] = None,
+        exc_val: typing.Optional[BaseException] = None,
+        exc_tb: typing.Optional[TracebackType] = None,
+    ) -> None:
+        await super().__aexit__(exc_type, exc_val, exc_tb)
+        await self.exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def handle_async_request(self, request: Request) -> Response:
+        scheme = request.url.scheme
+        headers = request.headers
+
+        if scheme in {"ws", "wss"} or headers.get("upgrade") == "websocket":
+            subprotocols: list[str] = []
+            if (
+                subprotocols_header := headers.get("sec-websocket-protocol")
+            ) is not None:
+                subprotocols = subprotocols_header.split(",")
+
+            scope = {
+                "type": "websocket",
+                "path": request.url.path,
+                "raw_path": request.url.raw_path,
+                "root_path": self.root_path,
+                "scheme": scheme,
+                "query_string": request.url.query,
+                "headers": [(k.lower(), v) for (k, v) in request.headers.raw],
+                "client": self.client,
+                "server": (request.url.host, request.url.port),
+                "subprotocols": subprotocols,
+            }
+            return await self._handle_ws_request(request, scope)
+
+        return await super().handle_async_request(request)
+
+    async def _handle_ws_request(
+        self,
+        request: Request,
+        scope: Scope,
+    ) -> Response:
+        assert isinstance(request.stream, AsyncByteStream)
+
+        self.scope = scope
+        async with contextlib.AsyncExitStack() as stack:
+            stream, accept_response = await stack.enter_async_context(
+                ASGIWebSocketAsyncNetworkStream(self.app, self.scope, self._task_group)  # type: ignore[arg-type]
+            )
+            self.exit_stacks.append(stack.pop_all())
+
+        accept_response_lines = accept_response.decode("utf-8").splitlines()
+        headers = [
+            typing.cast(tuple[str, str], line.split(": ", 1))
+            for line in accept_response_lines[1:]
+            if line.strip() != ""
+        ]
+
+        return Response(
+            status_code=101,
+            headers=headers,
+            extensions={"network_stream": stream},
+        )
+
+    async def aclose(self) -> None:
+        for stack in self.exit_stacks[::-1]:
+            await stack.aclose()