Added output index tracker

Author: 3coins

jupyter_server_documents/outputs/output_index_tracker.py

@@ -0,0 +1,88 @@
+import sys
+from typing import Dict, Set, Optional
+
+class OutputIndexTracker:
+    __slots__ = [
+        '_last_output_index', 
+        '_display_id_to_output_index', 
+        '_cell_display_ids'
+    ]
+    
+    def __init__(self):
+        # Dictionaries to track indices and display IDs
+        self._last_output_index: Dict[str, int] = {}
+        self._display_id_to_output_index: Dict[str, int] = {}
+        self._cell_display_ids: Dict[str, Set[str]] = {}
+    
+    def _intern_key(self, key: str) -> str:
+        """
+        Intern string keys to reduce memory usage.
+        Only use for keys that are likely to be repeated.
+        """
+        return sys.intern(key) if key else key
+    
+    def allocate_output_index(
+        self, 
+        cell_id: str, 
+        display_id: Optional[str] = None
+    ) -> int:
+        # Intern keys to reduce memory
+        cell_id = self._intern_key(cell_id)
+        display_id = self._intern_key(display_id) if display_id else None
+        
+        # Retrieve the last index for this cell, defaulting to -1 if not found
+        last_index = self._last_output_index.get(cell_id, -1)
+        
+        # If a display_id is provided, check for existing index
+        if display_id:
+            # If display_id already has an index, return that
+            if display_id in self._display_id_to_output_index:
+                return self._display_id_to_output_index[display_id]
+            
+            # Allocate a new index for this display_id
+            new_index = last_index + 1
+            self._display_id_to_output_index[display_id] = new_index
+            
+            # Track display_id for this cell
+            if cell_id not in self._cell_display_ids:
+                self._cell_display_ids[cell_id] = set()
+            self._cell_display_ids[cell_id].add(display_id)
+        else:
+            # For non-display specific outputs, simply increment the last index
+            new_index = last_index + 1
+        
+        # Update the last output index for the cell
+        self._last_output_index[cell_id] = new_index
+        
+        return new_index
+    
+    def get_output_index(self, display_id: str) -> Optional[int]:
+        """
+        Retrieve the output index for a given display ID.
+        
+        Args:
+            display_id (str): The display identifier.
+        
+        Returns:
+            Optional[int]: The output index if found, None otherwise.
+        """
+        return self._display_id_to_output_index.get(display_id)
+    
+    def clear_cell_indices(self, cell_id: str) -> None:
+        """
+        Clear indices associated with a specific cell.
+        
+        Args:
+            cell_id (str): The identifier of the cell to clear.
+        """
+        # Remove the last output index for this cell
+        self._last_output_index.pop(cell_id, None)
+        
+        # Remove and clean up associated display IDs
+        if cell_id in self._cell_display_ids:
+            # Remove all display IDs associated with this cell from the output index mapping
+            for display_id in self._cell_display_ids[cell_id]:
+                self._display_id_to_output_index.pop(display_id, None)
+            
+            # Remove the cell's display ID tracking
+            del self._cell_display_ids[cell_id]

jupyter_server_documents/tests/test_output_index_tracker.py

@@ -0,0 +1,119 @@
+import pytest
+import sys
+from jupyter_server_documents.outputs.output_index_tracker import OutputIndexTracker
+
+def test_basic_output_index_allocation():
+    """
+    Test basic output index allocation for a cell without display ID
+    """
+    tracker = OutputIndexTracker()
+    
+    # First output for a cell should be 0
+    assert tracker.allocate_output_index('cell1') == 0
+    assert tracker.allocate_output_index('cell1') == 1
+    assert tracker.allocate_output_index('cell1') == 2
+
+def test_output_index_with_display_id():
+    """
+    Test output index allocation with display IDs
+    """
+    tracker = OutputIndexTracker()
+    
+    # First output for a cell with display ID
+    assert tracker.allocate_output_index('cell1', 'display1') == 0
+    
+    # Subsequent calls with same display ID should return the same index
+    assert tracker.allocate_output_index('cell1', 'display1') == 0
+    
+    # Different display ID should get a new index
+    assert tracker.allocate_output_index('cell1', 'display2') == 1
+
+def test_multiple_cells_output_indices():
+    """
+    Test output index allocation across multiple cells
+    """
+    tracker = OutputIndexTracker()
+    
+    assert tracker.allocate_output_index('cell1') == 0
+    assert tracker.allocate_output_index('cell1') == 1
+    assert tracker.allocate_output_index('cell2') == 0
+    assert tracker.allocate_output_index('cell2') == 1
+
+def test_display_id_index_retrieval():
+    """
+    Test retrieving output index for a display ID
+    """
+    tracker = OutputIndexTracker()
+    
+    tracker.allocate_output_index('cell1', 'display1')
+    
+    assert tracker.get_output_index('display1') == 0
+    assert tracker.get_output_index('non_existent_display') is None
+
+def test_clear_cell_indices():
+    """
+    Test clearing indices for a specific cell
+    """
+    tracker = OutputIndexTracker()
+    
+    # Allocate some indices
+    tracker.allocate_output_index('cell1', 'display1')
+    tracker.allocate_output_index('cell1', 'display2')
+    tracker.allocate_output_index('cell2')
+    
+    # Clear cell1 indices
+    tracker.clear_cell_indices('cell1')
+    
+    # Verify cell1 indices are cleared
+    assert 'cell1' not in tracker._last_output_index
+    assert 'display1' not in tracker._display_id_to_output_index
+    assert 'display2' not in tracker._display_id_to_output_index
+    
+    # Verify cell2 indices remain
+    assert 'cell2' in tracker._last_output_index
+
+def test_memory_efficiency():
+    """
+    Basic test to ensure __slots__ is working
+    """
+    tracker = OutputIndexTracker()
+    
+    # Attempt to add a new attribute should raise an AttributeError
+    with pytest.raises(AttributeError):
+        tracker.new_attribute = "test"
+
+def test_key_interning():
+    """
+    Test that key interning works for repeated strings
+    """
+    tracker = OutputIndexTracker()
+    
+    # Use the same string multiple times
+    cell_id = "repeated_cell"
+    display_id = "repeated_display"
+    
+    # Allocate indices multiple times
+    for _ in range(5):
+        index1 = tracker.allocate_output_index(cell_id, display_id)
+        index2 = tracker.allocate_output_index(cell_id, display_id)
+        assert index1 == index2
+
+def test_cell_display_ids_tracking():
+    """
+    Test tracking of display IDs for a cell
+    """
+    tracker = OutputIndexTracker()
+    
+    # Allocate multiple display IDs for a cell
+    tracker.allocate_output_index('cell1', 'display1')
+    tracker.allocate_output_index('cell1', 'display2')
+    
+    # Verify display IDs are tracked
+    assert 'cell1' in tracker._cell_display_ids
+    assert set(tracker._cell_display_ids['cell1']) == {'display1', 'display2'}
+    
+    # Clear cell indices
+    tracker.clear_cell_indices('cell1')
+    
+    # Verify display IDs are cleared
+    assert 'cell1' not in tracker._cell_display_ids