feat: add alternative method and choose method

tests/test_stream.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Union
+from typing import List, Tuple, Union
 
 import pytest
 
@@ -163,3 +163,72 @@ def argument(stream: TokenStream) -> Union[str, Tuple[int, int, int]]:
             (1, 2, 3),
             "thing",
         ]
+
+
+def test_alternative():
+    stream = TokenStream("hello world 1 2 3 thing")
+
+    def argument(stream: TokenStream) -> Union[str, Tuple[int, int, int]]:
+        with stream.alternative():
+            return (
+                int(stream.expect("number").value),
+                int(stream.expect("number").value),
+                int(stream.expect("number").value),
+            )
+        return stream.expect("word").value  # type: ignore
+
+    with stream.syntax(number=r"\d+", word=r"\w+"):
+        assert [argument(stream) for _ in stream.peek_until()] == [
+            "hello",
+            "world",
+            (1, 2, 3),
+            "thing",
+        ]
+
+
+def test_choose():
+    stream = TokenStream("hello world 1 2 3 thing")
+
+    def word(stream: TokenStream) -> str:
+        return stream.expect("word").value
+
+    def triplet(stream: TokenStream) -> Tuple[int, int, int]:
+        return (
+            int(stream.expect("number").value),
+            int(stream.expect("number").value),
+            int(stream.expect("number").value),
+        )
+
+    def argument(stream: TokenStream) -> Union[str, Tuple[int, int, int]]:  # type: ignore
+        for parser, alternative in stream.choose(word, triplet):
+            with alternative:
+                return parser(stream)
+
+    with stream.syntax(number=r"\d+", word=r"\w+"):
+        assert [argument(stream) for _ in stream.peek_until()] == [
+            "hello",
+            "world",
+            (1, 2, 3),
+            "thing",
+        ]
+
+
+def test_choose_append():
+    stream = TokenStream("hello world 1 2 3 thing")
+    result: List[Union[str, Tuple[int, int, int]]] = []
+
+    with stream.syntax(number=r"\d+", word=r"\w+"):
+        while stream.peek():
+            for argument_type, alternative in stream.choose("word", "triplet"):
+                with alternative:
+                    result.append(
+                        stream.expect("word").value
+                        if argument_type == "word"
+                        else (
+                            int(stream.expect("number").value),
+                            int(stream.expect("number").value),
+                            int(stream.expect("number").value),
+                        )
+                    )
+
+    assert result == ["hello", "world", (1, 2, 3), "thing"]

tokenstream/stream.py

@@ -5,25 +5,30 @@
 ]
 
 import re
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass, field
 from typing import (
     Any,
     Callable,
     ClassVar,
+    ContextManager,
     Dict,
     Iterable,
     Iterator,
     List,
     Optional,
     Set,
     Tuple,
+    TypeVar,
     overload,
 )
 
 from .error import InvalidSyntax, UnexpectedEOF, UnexpectedToken
 from .token import SourceLocation, Token, TokenPattern
 
+T = TypeVar("T")
+
+
 SyntaxRules = Tuple[Tuple[str, str], ...]
 CheckpointCommit = Callable[[], None]
 
@@ -862,3 +867,49 @@ def checkpoint(self) -> Iterator[CheckpointCommit]:
         finally:
             if previous_index:
                 self.index = previous_index[0]
+
+    @contextmanager
+    def alternative(self) -> Iterator[None]:
+        """Keep going if the code within the ``with`` statement raises a syntax error.
+
+        >>> stream = TokenStream("hello world 123")
+        >>> with stream.syntax(word=r"[a-z]+", number=r"[0-9]+"):
+        ...     stream.expect("word").value
+        ...     stream.expect("word").value
+        ...     with stream.alternative():
+        ...         stream.expect("word").value
+        ...     stream.expect("number").value
+        'hello'
+        'world'
+        '123'
+        """
+        with self.checkpoint() as commit:
+            yield
+            commit()
+
+    def choose(self, *args: T) -> Iterator[Tuple[T, ContextManager[None]]]:
+        """Iterate over each argument until one of the alternative succeeds.
+
+        >>> stream = TokenStream("hello world 123")
+        >>> with stream.syntax(word=r"[a-z]+", number=r"[0-9]+"):
+        ...     while stream.peek():
+        ...         for token_type, alternative in stream.choose("word", "number"):
+        ...             with alternative:
+        ...                 stream.expect(token_type).value
+        'hello'
+        'world'
+        '123'
+        """
+        should_break = False
+
+        @contextmanager
+        def alternative():
+            with self.alternative():
+                nonlocal should_break
+                yield
+                should_break = True
+
+        for i, arg in enumerate(args):
+            yield arg, nullcontext() if i == len(args) - 1 else alternative()
+            if should_break:
+                break