Corrected thee Transformer's whole_tree interface, for both internal and external use

lark/parse_tree_builder.py

@@ -199,8 +199,7 @@ def inplace_transformer(func):
     def f(children):
         # function name in a Transformer is a rule name.
         tree = Tree(func.__name__, children)
-        func(tree)
-        return tree
+        return func(tree)
     return f
 
 class ParseTreeBuilder:

lark/visitors.py

@@ -36,7 +36,7 @@ class Transformer:
                     return f(*children)
                 elif getattr(f, 'whole_tree', False):
                     if new_children is not None:
-                        raise NotImplementedError("Doesn't work with the base Transformer class")
+                        tree.children = new_children
                     return f(tree)
                 else:
                     return f(children)

tests/test_parser.py

@@ -151,30 +151,50 @@ class TestParsers(unittest.TestCase):
         self.assertEqual( r.children, ["<c>"] )
 
     def test_embedded_transformer_inplace(self):
+        @v_args(tree=True)
         class T1(Transformer_InPlace):
             def a(self, tree):
-                assert isinstance(tree, Tree)
+                assert isinstance(tree, Tree), tree
                 tree.children.append("tested")
+                return tree
+
+            def b(self, tree):
+                return Tree(tree.data, tree.children + ['tested2'])
 
         @v_args(tree=True)
         class T2(Transformer):
             def a(self, tree):
                 assert isinstance(tree, Tree)
                 tree.children.append("tested")
+                return tree
+
+            def b(self, tree):
+                return Tree(tree.data, tree.children + ['tested2'])
 
         class T3(Transformer):
             @v_args(tree=True)
             def a(self, tree):
                 assert isinstance(tree, Tree)
                 tree.children.append("tested")
+                return tree
+
+            @v_args(tree=True)
+            def b(self, tree):
+                return Tree(tree.data, tree.children + ['tested2'])
 
         for t in [T1(), T2(), T3()]:
-            g = Lark("""start: a
-                        a : "x"
-                     """, parser='lalr', transformer=t)
-            r = g.parse("x")
-            first, = r.children
-            self.assertEqual(first.children, ["tested"])
+            for internal in [False, True]:
+                g = Lark("""start: a b
+                            a : "x"
+                            b : "y"
+                        """, parser='lalr', transformer=t if internal else None)
+                r = g.parse("xy")
+                if not internal:
+                    r = t.transform(r)
+
+                a, b = r.children
+                self.assertEqual(a.children, ["tested"])
+                self.assertEqual(b.children, ["tested2"])
 
     def test_alias(self):
         Lark("""start: ["a"] "b" ["c"] "e" ["f"] ["g"] ["h"] "x" -> d """)