|
3 | 3 | from __future__ import annotations
|
4 | 4 |
|
5 | 5 | import asyncio
|
| 6 | +from typing import Literal |
6 | 7 |
|
7 | 8 | from jupyter_server.extension.application import ExtensionApp
|
| 9 | +from jupyter_ydoc import ydocs as YDOCS |
| 10 | +from jupyter_ydoc.ybasedoc import YBaseDoc |
| 11 | +from pycrdt import Doc |
8 | 12 | from pycrdt_websocket.ystore import BaseYStore
|
9 | 13 | from traitlets import Bool, Float, Type
|
10 | 14 |
|
11 | 15 | from .handlers import DocSessionHandler, YDocWebSocketHandler
|
12 | 16 | from .loaders import FileLoaderMapping
|
| 17 | +from .rooms import DocumentRoom |
13 | 18 | from .stores import SQLiteYStore
|
14 |
| -from .utils import AWARENESS_EVENTS_SCHEMA_PATH, EVENTS_SCHEMA_PATH |
15 |
| -from .websocketserver import JupyterWebsocketServer |
| 19 | +from .utils import ( |
| 20 | + AWARENESS_EVENTS_SCHEMA_PATH, |
| 21 | + EVENTS_SCHEMA_PATH, |
| 22 | + encode_file_path, |
| 23 | + room_id_from_encoded_path, |
| 24 | +) |
| 25 | +from .websocketserver import JupyterWebsocketServer, RoomNotFound |
16 | 26 |
|
17 | 27 |
|
18 | 28 | class YDocExtension(ExtensionApp):
|
@@ -124,6 +134,43 @@ def initialize_handlers(self):
|
124 | 134 | ]
|
125 | 135 | )
|
126 | 136 |
|
| 137 | + async def get_document( |
| 138 | + self: YDocExtension, |
| 139 | + *, |
| 140 | + path: str, |
| 141 | + content_type: Literal["notebook", "file"], |
| 142 | + file_format: Literal["json", "text"], |
| 143 | + copy: bool = True, |
| 144 | + ) -> YBaseDoc | None: |
| 145 | + """Get a view of the shared model for the matching document. |
| 146 | +
|
| 147 | + If `copy=True`, the returned shared model is a fork, meaning that any changes |
| 148 | + made to it will not be propagated to the shared model used by the application. |
| 149 | + """ |
| 150 | + file_id_manager = self.serverapp.web_app.settings["file_id_manager"] |
| 151 | + file_id = file_id_manager.index(path) |
| 152 | + |
| 153 | + encoded_path = encode_file_path(file_format, content_type, file_id) |
| 154 | + room_id = room_id_from_encoded_path(encoded_path) |
| 155 | + |
| 156 | + try: |
| 157 | + room = await self.ywebsocket_server.get_room(room_id) |
| 158 | + except RoomNotFound: |
| 159 | + return None |
| 160 | + |
| 161 | + if isinstance(room, DocumentRoom): |
| 162 | + if copy: |
| 163 | + update = room.ydoc.get_update() |
| 164 | + |
| 165 | + fork_ydoc = Doc() |
| 166 | + fork_ydoc.apply_update(update) |
| 167 | + |
| 168 | + return YDOCS.get(content_type, YDOCS["file"])(fork_ydoc) |
| 169 | + else: |
| 170 | + return room._document |
| 171 | + |
| 172 | + return None |
| 173 | + |
127 | 174 | async def stop_extension(self):
|
128 | 175 | # Cancel tasks and clean up
|
129 | 176 | await asyncio.wait(
|
|
0 commit comments