Add merge function

lark/visitors.py

@@ -149,6 +149,62 @@ class Transformer(_Decoratable):
         return token
 
 
+def merge_transformers(base_transformer=None, **kwargs):
+    """
+    Add the methods of other transformer to this one.
+
+    This method is meant to aid in the maintenance of imports.
+
+    Example:
+    ```python
+    class T1(Transformer):
+        def method_on_main_grammar(self, children):
+            # Do something
+            pass
+
+        def imported_grammar__method(self, children):
+            # Do something else
+            pass
+
+    class TMain(Transformer):
+        def method_on_main_grammar(self, children):
+            # Do something
+            pass
+
+    class TImportedGrammar(Transformer):
+        def method(self, children):
+            # Do something else
+            pass
+
+    regular_transformer = T1()
+    composed_transformer = merge_transformers(TMain(),
+                                              imported_grammar=TImportedGrammar())
+    ```
+    In the above code block `regular_transformer` and `composed_transformer`
+    should behave identically.
+    """
+
+    if base_transformer is None:
+        base_transformer = Transformer()
+    for prefix, transformer in kwargs.items():
+        prefix += "__"
+
+        for method_name in dir(transformer):
+            method = getattr(transformer, method_name)
+            if not callable(method):
+                continue
+            if method_name.startswith("_") or method_name == "transform":
+                continue
+            new_method_name = prefix + method_name
+            if prefix + method_name in dir(base_transformer):
+                raise AttributeError(
+                    ("Method '{new_method_name}' already present in base "
+                     "transformer").format(new_method_name=new_method_name))
+            setattr(base_transformer, new_method_name, method)
+
+    return base_transformer
+
+
 class InlineTransformer(Transformer):   # XXX Deprecated
     def _call_userfunc(self, tree, new_children=None):
         # Assumes tree is already transformed

tests/test_trees.py

@@ -9,7 +9,7 @@ import functools
 from lark.tree import Tree
 from lark.lexer import Token
 from lark.visitors import Visitor, Visitor_Recursive, Transformer, Interpreter, visit_children_decor, v_args, Discard, Transformer_InPlace, \
-    Transformer_InPlaceRecursive, Transformer_NonRecursive
+    Transformer_InPlaceRecursive, Transformer_NonRecursive, merge_transformers
 
 
 class TestTrees(TestCase):
@@ -233,21 +233,69 @@ class TestTrees(TestCase):
 
         x = MyTransformer().transform( t )
         self.assertEqual(x, t2)
-    
+
     def test_transformer_variants(self):
         tree = Tree('start', [Tree('add', [Token('N', '1'), Token('N', '2')]), Tree('add', [Token('N', '3'), Token('N', '4')])])
         for base in (Transformer, Transformer_InPlace, Transformer_NonRecursive, Transformer_InPlaceRecursive):
             class T(base):
                 def add(self, children):
                     return sum(children)
-                
+
                 def N(self, token):
                     return int(token)
-            
+
             copied = copy.deepcopy(tree)
             result = T().transform(copied)
             self.assertEqual(result, Tree('start', [3, 7]))
 
+    def test_merge_transformers(self):
+        tree = Tree('start', [
+            Tree('main', [
+                Token("A", '1'), Token("B", '2')
+            ]),
+            Tree("module__main", [
+                Token("A", "2"), Token("B", "3")
+            ])
+        ])
+
+        class T1(Transformer):
+            A = int
+            B = int
+            main = sum
+            start = list
+            def module__main(self, children):
+                prod = 1
+                for child in children:
+                    prod *= child
+                return prod
+
+        class T2(Transformer):
+            A = int
+            B = int
+            main = sum
+            start = list
+
+        class T3(Transformer):
+            def main(self, children):
+                prod = 1
+                for child in children:
+                    prod *= child
+                return prod
+
+        class T4(Transformer):
+            def other_aspect(self, children):
+                pass
+
+        t1_res = T1().transform(tree)
+        composed_res = merge_transformers(T2(), module=T3()).transform(tree)
+        self.assertEqual(t1_res, composed_res)
+        with self.assertRaises(AttributeError):
+            merge_transformers(T1(), module=T3())
+
+        try:
+            composed = merge_transformers(T1(), module=T4())
+        except AttributeError:
+            self.fail("Should be able to add classes that do not conflict")
 
 if __name__ == '__main__':
     unittest.main()