Generic tree typing

lark/visitors.py

@@ -1,4 +1,5 @@
 from typing import TypeVar, Tuple, List, Callable, Generic, Type, Union, Optional, cast
+
 from abc import ABC
 from functools import wraps
 
@@ -9,9 +10,14 @@
 
 ###{standalone
 from inspect import getmembers, getmro
-
-_Return_T = TypeVar('_Return_T')
-_Return_V = TypeVar('_Return_V')
+import sys
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+else:
+    from typing_extensions import Protocol
+
+_Return_T = TypeVar('_Return_T', covariant=True)
+_Return_V = TypeVar('_Return_V', covariant=True)
 _Leaf_T = TypeVar('_Leaf_T')
 _Leaf_U = TypeVar('_Leaf_U')
 _R = TypeVar('_R')
@@ -55,6 +61,21 @@ def __class_getitem__(cls, _):
         return cls
 
 
+class ITransformer(Protocol[_Leaf_T, _Return_T]):
+    """Base interface for classes that can transform a Tree"""
+
+    def transform(self, tree: Tree[_Leaf_T]) -> _Return_T:
+        """Transform the given tree, and return the final result"""
+        pass
+
+    def __mul__(
+            self: 'ITransformer[_Leaf_T, Tree[_Leaf_U]]',
+            other: 'ITransformer[_Leaf_U, _Return_V]'
+    ) -> 'ITransformer[_Leaf_T, _Return_V]':
+        """Chain two transformers together, returning a new transformer."""
+        pass
+
+
 class Transformer(_Decoratable, ABC, Generic[_Leaf_T, _Return_T]):
     """Transformers visit each node of the tree, and run the appropriate method on it according to the node's data.
 
@@ -140,8 +161,8 @@ def transform(self, tree: Tree[_Leaf_T]) -> _Return_T:
 
     def __mul__(
             self: 'Transformer[_Leaf_T, Tree[_Leaf_U]]',
-            other: 'Union[Transformer[_Leaf_U, _Return_V], TransformerChain[_Leaf_U, _Return_V,]]'
-    ) -> 'TransformerChain[_Leaf_T, _Return_V]':
+            other: ITransformer[_Leaf_U, _Return_V]
+    ) -> ITransformer[_Leaf_T, _Return_V]:
         """Chain two transformers together, returning a new transformer.
         """
         return TransformerChain(self, other)
@@ -163,9 +184,9 @@ def __default_token__(self, token):
 
 class TransformerChain(Generic[_Leaf_T, _Return_T]):
 
-    transformers: 'Tuple[Union[Transformer, TransformerChain], ...]'
+    transformers: Tuple[ITransformer, ...]
 
-    def __init__(self, *transformers: 'Union[Transformer, TransformerChain]') -> None:
+    def __init__(self, *transformers: ITransformer) -> None:
         self.transformers = transformers
 
     def transform(self, tree: Tree[_Leaf_T]) -> _Return_T:
@@ -175,8 +196,8 @@ def transform(self, tree: Tree[_Leaf_T]) -> _Return_T:
 
     def __mul__(
             self: 'TransformerChain[_Leaf_T, Tree[_Leaf_U]]',
-            other: 'Union[Transformer[_Leaf_U, _Return_V], TransformerChain[_Leaf_U, _Return_V]]'
-    ) -> 'TransformerChain[_Leaf_T, _Return_V]':
+            other: ITransformer[_Leaf_U, _Return_V]
+    ) -> ITransformer[_Leaf_T, _Return_V]:
         return TransformerChain(*self.transformers + (other,))
 
 

setup.py

@@ -12,7 +12,7 @@
     packages = ['lark', 'lark.parsers', 'lark.tools', 'lark.grammars', 'lark.__pyinstaller'],
 
     requires = [],
-    install_requires = [],
+    install_requires = ['typing_extensions >= 3.7, < 5.0; python_version < "3.8"'],
 
     extras_require = {
         "regex": ["regex"],

test-requirements.txt

@@ -1,2 +1,3 @@
 Js2Py==0.68
-regex
+regex
+typing_extensions >= 3.7, < 5.0; python_version < "3.8"