feat: add token_handler

Author: vberlier

tokenstream/stream.py

@@ -8,7 +8,7 @@
 import re
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Any, ContextManager, Iterable, Iterator, TypeVar, overload
+from typing import Any, Callable, ContextManager, Iterable, Iterator, TypeVar, overload
 
 from .error import InvalidSyntax, UnexpectedEOF, UnexpectedToken
 from .location import SourceLocation, set_location
@@ -68,6 +68,9 @@ class TokenStream:
         >>> stream.source
         'hello world'
 
+    token_handler
+        A callback for modifying tokens before they're emitted.
+
     syntax_rules
         A tuple of ``(token_type, pattern)`` pairs that define the recognizable tokens.
 
@@ -140,6 +143,7 @@ class TokenStream:
     """
 
     source: str
+    token_handler: Callable[[Token], Token] | None = extra_field(default=None)
     syntax_rules: SyntaxRules = extra_field(default=())
     regex: re.Pattern[str] = extra_field()
 
@@ -482,6 +486,9 @@ def emit_token(self, token_type: str, value: str = "") -> Token:
             end_location=SourceLocation(end_pos, end_lineno, end_colno),
         )
 
+        if self.token_handler:
+            token = self.token_handler(token)
+
         self.location = token.end_location
         self.tokens.append(token)