diff --git a/advent/parser/parser.py b/advent/parser/parser.py index 1b5edba..3398307 100644 --- a/advent/parser/parser.py +++ b/advent/parser/parser.py @@ -127,6 +127,18 @@ class P(Generic[T]): def optional(self) -> P[T | None]: return P.either(self, P.pure(None)) + def times(self, *, max: int | None = None, min: int | None = None, + exact: int | None = None) -> P[list[T]]: + match (exact, min, max): + case (int(e), None, None): + return self.many().satisfies(lambda lst: len(lst) == e) + case (None, int(mn), None): + return self.many().satisfies(lambda lst: len(lst) >= mn) + case (None, None, int(mx)): + return self.many().satisfies(lambda lst: len(lst) <= mx) + case _: + raise Exception("Choose exactly one of exact, min or max") + def sep_by(self, sep: P[Any]) -> P[list[T]]: return P.map2(self, P.snd(sep, self).many(), lambda f, r: [f] + r) diff --git a/advent/parser/test_parser.py b/advent/parser/test_parser.py index d119444..61d6a2c 100644 --- a/advent/parser/test_parser.py +++ b/advent/parser/test_parser.py @@ -181,3 +181,27 @@ def test_choice(): expected = ['1', 1] result = list(parser.parse_multi(input)) assert result == expected + + +def test_times_exact(): + input = 'aaa' + parser = P.is_char('a').times(exact=2) + expected = [['a', 'a']] + result = list(parser.parse_multi(input)) + assert result == expected + + +def test_times_min(): + input = 'aaa' + parser = P.is_char('a').times(min=2) + expected = [['a', 'a', 'a'], ['a', 'a']] + result = list(parser.parse_multi(input)) + assert result == expected + + +def test_times_max(): + input = 'aaa' + parser = P.is_char('a').times(max=2) + expected = [['a', 'a'], ['a'], []] + result = list(parser.parse_multi(input)) + assert result == expected