Added a state to track display_ids, added tests

Author: 3coins

jupyter_server_documents/outputs/manager.py

@@ -13,8 +13,9 @@
 
 class OutputsManager(LoggingConfigurable):
     _last_output_index = Dict(default_value={})
+    _display_id_output_index = Dict(default_value={})
+    _display_ids = Dict(default_value={})
     _stream_count = Dict(default_value={})
-    _output_index_by_display_id = Dict(default_value={})
 
     outputs_path = Instance(PurePath, help="The local runtime dir")
     stream_limit = Int(default_value=10, config=True, allow_none=True)
@@ -35,24 +36,28 @@ def _build_path(self, file_id, cell_id=None, output_index=None):
             path = path / f"{output_index}.output"
         return path
 
-    def _determine_output_index(self, cell_id, display_id=None):
+    def _compute_output_index(self, cell_id, display_id=None):
         """
-        Determines output index for a given cell and optional display ID.
+        Computes next output index for a cell.
         
         Args:
             cell_id (str): The cell identifier
             display_id (str, optional): A display identifier. Defaults to None.
         
         Returns:
-            int: The allocated output index
+            int: The output index
         """
         last_index = self._last_output_index.get(cell_id, -1)
         if display_id:
-            index = self._output_index_by_display_id.get(display_id)
+            if cell_id not in self._display_ids:
+                self._display_ids[cell_id] = set([display_id])
+            else:
+                self._display_ids[cell_id].add(display_id)
+            index = self._display_id_output_index.get(display_id)
             if index is None:
                 index = last_index + 1
                 self._last_output_index[cell_id] = index
-                self._output_index_by_display_id[display_id] = index
+                self._display_id_output_index[display_id] = index
         else:
             index = last_index + 1
             self._last_output_index[cell_id] = index
@@ -61,7 +66,7 @@ def _determine_output_index(self, cell_id, display_id=None):
 
     def get_output_index(self, display_id: str):
         """Returns output index for a cell by display_id"""
-        return self._output_index_by_display_id.get(display_id)
+        return self._display_id_output_index.get(display_id)
 
     def get_output(self, file_id, cell_id, output_index):
         """Get an output by file_id, cell_id, and output_index."""
@@ -120,7 +125,7 @@ def write(self, file_id, cell_id, output, display_id=None):
 
     def write_output(self, file_id, cell_id, output, display_id=None):
         self._ensure_path(file_id, cell_id)
-        index = self._determine_output_index(cell_id, display_id)
+        index = self._compute_output_index(cell_id, display_id)
         path = self._build_path(file_id, cell_id, index)
         data = json.dumps(output, ensure_ascii=False)
         with open(path, "w", encoding="utf-8") as f:
@@ -161,17 +166,15 @@ def clear(self, file_id, cell_id=None):
         """Clear the state of the manager."""
         if cell_id is None:
             self._stream_count = {}
-            path = self._build_path(file_id)
         else:
-            try:
-                del self._stream_count[cell_id]
-            except KeyError:
-                pass
-            try:
-                del self._last_output_index[cell_id]
-            except KeyError:
-                pass
-            path = self._build_path(file_id, cell_id)
+            self._stream_count.pop(cell_id, None)
+            self._last_output_index.pop(cell_id, None)
+            
+            display_ids = self._display_ids.get(cell_id, [])
+            for display_id in display_ids:
+                self._display_id_output_index.pop(display_id, None)
+
+        path = self._build_path(file_id, cell_id)    
         try:
             shutil.rmtree(path)
         except FileNotFoundError:

jupyter_server_documents/tests/test_outputs_manager.py

@@ -94,39 +94,72 @@ def file_not_found():
             op.get_stream('a','b')       
 
 
-@pytest.mark.parametrize(
-    "_last_output_index, _output_index_by_display_id, cell_id, display_id, expected", [
-    # Empty dictionaries, no display_id
-    ({}, {}, "cell1", None, 0),
+def test__compute_output_index_basic():
+    """
+    Test basic output index allocation for a cell without display ID
+    """
+    op = OutputsManager()
 
-    # Empty dictionaries, with display_id
-    ({}, {}, "cell1", "display1", 0),
+    # First output for a cell should be 0
+    assert op._compute_output_index('cell1') == 0
+    assert op._compute_output_index('cell1') == 1
+    assert op._compute_output_index('cell1') == 2
+
+def test__compute_output_index_with_display_id():
+    """
+    Test output index allocation with display IDs
+    """
+    op = OutputsManager()
 
-    # Existing last_output_index, no display_id
-    ({"cell1": 5}, {}, "cell1", None, 6),
-    ({"cell2": 3}, {}, "cell1", None, 0),
+    # First output for a cell with display ID
+    assert op._compute_output_index('cell1', 'display1') == 0
 
-    # Existing output_index_by_display_id
-    ({}, {"display1": 2}, "cell1", "display1", 2),
+    # Subsequent calls with same display ID should return the same index
+    assert op._compute_output_index('cell1', 'display1') == 0
 
-    # Existing last_output_index and display_id
-    ({"cell1": 5}, {"display1": 3}, "cell1", "display1", 3),
-    ({"cell1": 5}, {"display1": 3}, "cell1", "display2", 6),
+    # Different display ID should get a new index
+    assert op._compute_output_index('cell1', 'display2') == 1
+
+
+def test__compute_output_index_multiple_cells():
+    """
+    Test output index allocation across multiple cells
+    """
+    op = OutputsManager()
 
-    # Multiple cells with different indices
-    ({"cell1": 2, "cell2": 7}, {}, "cell1", None, 3),
-    ({"cell1": 2, "cell2": 7}, {}, "cell3", None, 0),
-])
-def test__determine_output_index(
-    _last_output_index, 
-    _output_index_by_display_id, 
-    cell_id, 
-    display_id,
-    expected
-):
+    assert op._compute_output_index('cell1') == 0
+    assert op._compute_output_index('cell1') == 1
+    assert op._compute_output_index('cell2') == 0
+    assert op._compute_output_index('cell2') == 1
+
+def test_display_id_index_retrieval():
+    """
+    Test retrieving output index for a display ID
+    """
     op = OutputsManager()
-    op._last_output_index = _last_output_index;
-    op._output_index_by_display_id = _output_index_by_display_id
 
-    assert expected == op._determine_output_index(cell_id, display_id)
+    op._compute_output_index('cell1', 'display1')
+    
+    assert op.get_output_index('display1') == 0
+    assert op.get_output_index('non_existent_display') is None
 
+def test_display_ids():
+    """
+    Test tracking of display IDs for a cell
+    """
+    op = OutputsManager()
+    
+    # Allocate multiple display IDs for a cell
+    op._compute_output_index('cell1', 'display1')
+    op._compute_output_index('cell1', 'display2')
+    
+    # Verify display IDs are tracked
+    assert 'cell1' in op._display_ids
+    assert set(op._display_ids['cell1']) == {'display1', 'display2'}
+    
+    # Clear cell indices
+    op.clear('file1', 'cell1')
+    
+    # Verify display IDs are cleared
+    assert 'display1' not in op._display_ids
+    assert 'display2' not in op._display_ids