diff --git a/src/pick/__init__.py b/src/pick/__init__.py index 6858644..8c4713a 100644 --- a/src/pick/__init__.py +++ b/src/pick/__init__.py @@ -66,8 +66,11 @@ def __post_init__(self) -> None: if len(self.options) == 0: raise ValueError("options should not be an empty list") - if self.default_index >= len(self.options): - raise ValueError("default_index should be less than the length of options") + if self.default_index < -len(self.options) or self.default_index >= len(self.options): + raise ValueError("default_index is out of range of options") + + if self.default_index < 0: + self.default_index += len(self.options) if self.multiselect and self.min_selection_count > len(self.options): raise ValueError( diff --git a/tests/test_pick.py b/tests/test_pick.py index 759d086..b840333 100644 --- a/tests/test_pick.py +++ b/tests/test_pick.py @@ -1,4 +1,6 @@ -from pick import Picker, Option +import pytest + +from pick import Option, Picker def test_move_up_down(): @@ -78,6 +80,21 @@ def test_disabled_option(): picker.move_down() assert picker.get_selected() == (Option("option3"), 2) +def test_negative_default_index(): + options = ["option1", "option2", "option3"] + picker = Picker(options, default_index=-1) + assert picker.get_selected() == ("option3", 2) + lines, _ = picker.get_lines() + assert "* option3" in lines + + picker = Picker(options, default_index=-3) + assert picker.get_selected() == ("option1", 0) + + with pytest.raises(ValueError): + Picker(options, default_index=-4) + with pytest.raises(ValueError): + Picker(options, default_index=3) + def test_mark_index_disabled_option(): options = [Option("option1"), Option("option2", enabled=False), Option("option3")]