Fix test case

Author: Johannes E. M. Mosig

penzai/core/selectors.py

@@ -39,7 +39,7 @@
 T = typing.TypeVar("T")
 
 
-def shift_negative_indices(indices: Iterable[int], shift: int) -> tuple[int, ...]:
+def _shift_negative_indices(indices: Iterable[int], shift: int) -> tuple[int, ...]:
   """Adds `shift` to negative indices and leaves non-negative indices unchanged
 
   Can be used to handle negative indices. For example, if we expect indices in
@@ -49,12 +49,24 @@ def shift_negative_indices(indices: Iterable[int], shift: int) -> tuple[int, ...
   shift_negative_indices([0, 3, -2], len(r))
   ```
 
-  to get `(0, 3, 4)`
+  to get `(0, 3, 4)`. The same can be achieved in more generality with
+
+  ```py
+  pz.select((0, 3, -2)) \
+    .at_instances_of(int) \
+    .where(lambda i: i < 0) \
+    .apply(lambda i: i + shift)
+  ```
+
+  which is why this method is private to this module
 
   Args:
     indices: The integers to shift
     shift: The offset to add to negative indices. Usually, this is the largest
       index + 1, i.e. the length of the range of indices
+
+  Returns:
+    The indices as a tuple, with negative indices increased by `shift`
   """
   maybe_shifted_indices = []
   for index in indices:
@@ -1382,7 +1394,7 @@ def pick_nth_selected(self, n: int | Sequence[int]) -> Selection:
     else:
       indices = n
 
-    indices = shift_negative_indices(indices, len(self.selected_by_path))
+    indices = _shift_negative_indices(indices, len(self.selected_by_path))
 
     with _wrap_selection_errors(self):
       keep = _InProgressSelectionBoundary

tests/core/selectors_test.py

@@ -32,7 +32,7 @@
 @pytest.mark.parametrize(
   "input_indices, shift, expected_output",
   [
-    ((,), 1, (,)),
+    ((), 1, ()),
     ([0, 3, -2], len(range(6)), (0, 3, 4)),
   ]
 )
@@ -42,7 +42,7 @@ def test_shift_negative_indices(
   expected_output: tuple[int, ...],
 ):
   assert (
-    penzai.core.selectors.shift_negative_indices(input_indices, shift=shift)
+    penzai.core.selectors._shift_negative_indices(input_indices, shift=shift)
     == expected_output
   )