Repository: google-deepmind/penzai

Issue:
pick_nth_selected(-1) should select the last node
I would expect 
```
pz.select([0, 1, 2]).at_instances_of(int).pick_nth_selected(-1)
```
to return the same as 
```
pz.select([0, 1, 2]).at_instances_of(int).pick_nth_selected(2)
```
, but instead I get 

```
pz.select(
  [
    0,
    1,
    2,
  ]
).at_keypaths(<0 subtrees, highlighted above>)
```

It'd be great if `pick_nth_selected` would also support negative values, as I often want to select the last item.

Hints from issue comments:
@JEM-Mosig: I can't directly open pull requests, so I just leave it here as a comment. One way of implementing this could be to insert one line into the `pick_nth_selected` function:
```py
  def pick_nth_selected(self, n: int | Sequence[int]) -> Selection:
    """Filters a selection to only the nth selected node.

    ``n`` is taken relative to the linearized sequence of currently selected
    nodes, in the sense that ::

      my_selection.get_sequence()[n] == my_selection.pick_nth(n).get()

    Args:
      n: The index of the selected nodes that should remain selected. If this is
        a sequence, all nodes in the sequence will be selected.

    Returns:
      A new selection that includes only the ``n``th node from the original
      selection.  For instance, if the original selection has 5 nodes selected,
      setting ``n = 1`` would produce a new selection with only the node at
      index 1 selected.
    """
    if isinstance(n, int):
      indices = (n,)
    else:
      indices = n

    # NEW CODE
    # Handle negative indices
    indices = tuple(i % len(self.selected_by_path) for i in indices)
    # END NEW CODE

    with _wrap_selection_errors(self):
      keep = _InProgressSelectionBoundary
      new_selected_by_path = collections.OrderedDict({
          key: keep(value) if i in indices else value
          for i, (key, value) in enumerate(self.selected_by_path.items())
      })
      return _build_selection_from_boundary(
          Selection(
              selected_by_path=new_selected_by_path, remainder=self.remainder
          ).deselect()
      )
```
or, if performance is an issue, one could do the same mod-operation directly in the `if isinstance(n, int)` block@danieldjohnson: Hi, sorry for the delay on getting back to you on this. Agreed this seems like a useful thing to support!

Are you not able to open a pull request using the normal Github PR process? The repo is supposed to allow PRs. If you are running into problems, could you let me know what you are seeing on your end so that I can try to fix it?

Regarding the implementation, we probably should check for negative numbers separately instead of using %, since we wouldn't want `pz.select([0, 1, 2]).at_instances_of(int).pick_nth_selected(5)` to select the 2. Maybe adding `len(self.selected_by_path)` to negative numbers would do the trick.@JEM-Mosig: Hi @danieldjohnson . I tried again and discovered that you can do PRs from forks, so here it is now: https://github.com/google-deepmind/penzai/pull/110. 

How do you run tests locally? I haven't managed to get this to work (though I didn't try for long). Ideally it'd just be `uv run pytest`, but that doesn't work. (I saw the `run_tests.py` script, but did not manage to make this work either).@danieldjohnson: Thanks, just saw it!

The commands for running tests can be seen in the CI workflow script [here](https://github.com/google-deepmind/penzai/blob/main/.github/workflows/unittests.yml). Specifically:

```bash
# Lint / formatting:
uv run pyink penzai tests --check
uv run pylint penzai
# Unittests
uv run python run_tests.py
# Typechecking
uv run pytype --jobs auto penzai
```

Solution patch:
diff --git a/penzai/core/selectors.py b/penzai/core/selectors.py
index d89f8e0..3832fd9 100644
--- a/penzai/core/selectors.py
+++ b/penzai/core/selectors.py
@@ -39,6 +39,42 @@ OtherSubtree = typing.TypeVar("OtherSubtree")
 T = typing.TypeVar("T")
 
 
+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
+  `r = range(6)` and we get `[0, 3, -2]` as input, we can use ::
+
+    shift_negative_indices([0, 3, -2], len(r))
+
+  to get `(0, 3, 4)`. The same can be achieved in more generality with ::
+
+    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:
+    if index < 0:
+      maybe_shifted_indices.append(index + shift)
+    else:
+      maybe_shifted_indices.append(index)
+  return tuple(maybe_shifted_indices)
+
+
 @struct.pytree_dataclass
 class SelectionHole(struct.Struct):
   """A hole in a structure, taking the place of a selected subtree.
@@ -1356,6 +1392,8 @@ class Selection(Generic[SelectedSubtree], struct.Struct):
     else:
       indices = n
 
+    indices = _shift_negative_indices(indices, len(self.selected_by_path))
+
     with _wrap_selection_errors(self):
       keep = _InProgressSelectionBoundary
       new_selected_by_path = collections.OrderedDict({


Test patch:
diff --git a/tests/core/selectors_test.py b/tests/core/selectors_test.py
index 66eb1a2..f296ba0 100644
--- a/tests/core/selectors_test.py
+++ b/tests/core/selectors_test.py
@@ -56,6 +56,18 @@ class SELECTED_PART:  # pylint: disable=invalid-name
 
 class SelectorsTest(absltest.TestCase):
 
+  def test_shift_negative_indices(self):
+    for input_indices, shift, expected_output in [
+        ((), 1, ()),
+        ([0, 3, -2], len(range(6)), (0, 3, 4)),
+    ]:
+      self.assertEqual(
+          penzai.core.selectors._shift_negative_indices(
+              input_indices, shift=shift
+          ),
+          expected_output,
+      )
+
   def test_select(self):
     self.assertEqual(
         pz.select(make_example_object()),
@@ -565,6 +577,26 @@ class SelectorsTest(absltest.TestCase):
         ),
         [0, 1, 2, 3, SELECTED_PART(value=4), 5, 6, 7, 8, 9],
     )
+    # Test negative indices for `pick_nth_selected`
+    self.assertEqual(
+        (
+            pz.select(list(range(10)))
+            .at_instances_of(int)
+            .pick_nth_selected(-2)
+            .apply(SELECTED_PART)
+        ),
+        [0, 1, 2, 3, 4, 5, 6, 7, SELECTED_PART(value=8), 9],
+    )
+    # Don't select anything if index is out of range
+    self.assertEqual(
+        (
+            pz.select([0, 1, 2])
+            .at_instances_of(int)
+            .pick_nth_selected(5)
+            .apply(SELECTED_PART)
+        ),
+        [0, 1, 2],
+    )
 
   def test_invert__example_1(self):
     predicate = lambda node: isinstance(node, CustomLeaf) and node.tag <= 10