Better typing for combining multiple parsers

* Support for combining multiple parsers * Support for combining _TupleParser with Parser * Support for combining multiple _TupleParsers
vlasovskikh · Apr 16, 2024 · 9efde60 · 9efde60
1 parent dca8fd1
commit 9efde60
Showing 1 changed file with 32 additions and 7 deletions.
diff --git a/funcparserlib/parser.py b/funcparserlib/parser.py
@@ -66,6 +66,7 @@
 
 import logging
 import warnings
+import sys
 from typing import (
     Any,
     Callable,
@@ -82,6 +83,11 @@
 
 from funcparserlib.lexer import Token
 
+if sys.version_info >= (3, 11):
+    from typing import TypeVarTuple, Unpack
+else:
+    from typing_extensions import TypeVarTuple, Unpack
+
 log = logging.getLogger("funcparserlib")
 
 debug = False
@@ -273,13 +279,13 @@ def __add__(  # type: ignore[misc]
         pass
 
     @overload
-    def __add__(self, other: "Parser[_A, _C]") -> "_TupleParser[_A, Tuple[_B, _C]]":
+    def __add__(self, other: "Parser[_A, _C]") -> "_TupleParser[_A, _B, _C]":
         pass
 
     def __add__(
         self,
         other: Union["_IgnoredParser[_A]", "Parser[_A, _C]"],
-    ) -> Union["Parser[_A, _B]", "_TupleParser[_A, Tuple[_B, _C]]"]:
+    ) -> Union["Parser[_A, _B]", "_TupleParser[_A, _B, _C]"]:
         """Sequential combination of parsers. It runs this parser, then the other
         parser.
 
@@ -547,22 +553,41 @@ def __str__(self) -> str:
         return self.msg
 
 
+_Ts = TypeVarTuple("_Ts")
+_Ks = TypeVarTuple("_Ks")
+
+
 class _Tuple(tuple):
     pass
 
 
-class _TupleParser(Parser[_A, _B], Generic[_A, _B]):
+class _TupleParser(Parser[_A, Tuple[Unpack[_Ts]]], Generic[_A, Unpack[_Ts]]):
     @overload  # type: ignore[override]
-    def __add__(self, other: "_IgnoredParser[_A]") -> "_TupleParser[_A, _B]":
+    def __add__(self, other: "_IgnoredParser[_A]") -> "_TupleParser[_A, Unpack[_Ts]]":
         pass
 
     @overload
-    def __add__(self, other: Parser[_A, Any]) -> Parser[_A, Any]:
+    def __add__(
+        self, other: "_TupleParser[_A, Unpack[_Ks]]"
+    ) -> "_TupleParser[_A, Unpack[Tuple]]":
+        pass
+
+    @overload
+    def __add__(self, other: "Parser[_A, _B]") -> "_TupleParser[_A, Unpack[_Ts], _B]":
         pass
 
     def __add__(
-        self, other: Union["_IgnoredParser[_A]", Parser[_A, Any]]
-    ) -> Union["_TupleParser[_A, _B]", Parser[_A, Any]]:
+        self,
+        other: Union[
+            "_IgnoredParser[_A]",
+            Parser[_A, _B],
+            "_TupleParser[_A, Unpack[_Ks]]",
+        ],
+    ) -> Union[
+        "_TupleParser[_A, Unpack[_Ts]]",
+        "_TupleParser[_A, Unpack[_Ts], _B]",
+        "_TupleParser[_A, Unpack[Tuple]]",
+    ]:
         return super().__add__(other)