diff --git a/src/code/snap_operator.py b/src/code/snap_operator.py index e31c1ad..28b4fb8 100644 --- a/src/code/snap_operator.py +++ b/src/code/snap_operator.py @@ -57,6 +57,7 @@ def find_target(config: configs.Config, suffix: str) -> Optional[snap_holder.Sna def _all_but_last_k(array: list[_GenericT], k: int) -> Iterator[_GenericT]: + """All but at most k last elements.""" if k < 0: raise ValueError(f"k = {k} < 0") yield from array[: len(array) - k] diff --git a/src/code/snap_operator_test.py b/src/code/snap_operator_test.py index a50c693..6e9ad9e 100644 --- a/src/code/snap_operator_test.py +++ b/src/code/snap_operator_test.py @@ -140,8 +140,11 @@ def test_all_but_k(self): self.assertEqual( list(snap_operator._all_but_last_k([1, 2, 3, 4], 0)), [1, 2, 3, 4] ) + # Leave -2 out of 4 elements - causes an error. with self.assertRaisesRegex(ValueError, r"k = .+ < 0"): list(snap_operator._all_but_last_k([1, 2, 3, 4], -2)) + # Leave 5 out of 2 elements. Succeeds. + self.assertEqual(list(snap_operator._all_but_last_k([1, 2], 5)), []) def setUp(self) -> None: super().setUp()