pick_nth_selected(-1) should select the last node

[JEM-Mosig]

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.

[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:

  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: #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. Specifically:

# 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