From 607f984523fb5e062f53dee51dc68662b6a46979 Mon Sep 17 00:00:00 2001
From: Chanic Panic <chanicpanic@gmail.com>
Date: Fri, 28 Aug 2020 13:36:49 -0700
Subject: [PATCH] Create TreeForestTransformer

---
 lark/parsers/earley_forest.py | 81 ++++++++++++++++++++++++++++++++++-
 1 file changed, 80 insertions(+), 1 deletion(-)

diff --git a/lark/parsers/earley_forest.py b/lark/parsers/earley_forest.py
index a2f2ce2..1dfd3d1 100644
--- a/lark/parsers/earley_forest.py
+++ b/lark/parsers/earley_forest.py
@@ -12,8 +12,10 @@ from math import isinf
 from collections import deque
 from operator import attrgetter
 from importlib import import_module
+from functools import partial
 
-from .. visitors import Discard
+from ..parse_tree_builder import AmbiguousIntermediateExpander
+from ..visitors import Discard
 from ..lexer import Token
 from ..utils import logger
 from ..tree import Tree
@@ -497,6 +499,83 @@ class ForestToParseTree(ForestTransformer):
         self._on_cycle_retreat = False
         return super(ForestToParseTree, self).visit_token_node(node)
 
+def handles_ambiguity(func):
+    """Decorator for methods of subclasses of TreeForestTransformer.
+    Denotes that the method should receive a list of transformed derivations."""
+    func.handles_ambiguity = True
+    return func
+
+class TreeForestTransformer(ForestToParseTree):
+    """A ForestTransformer with a tree-Transformer-like interface.
+    By default, it will construct a tree.
+
+    Methods provided via inheritance are called based on the rule/symbol
+    names of nodes in the forest.
+
+    Methods that act on rules will receive a list of the results of the
+    transformations of the rule's children. By default, trees and tokens.
+
+    Methods that act on tokens will receive a Token.
+
+    Alternatively, methods that act on rules may be annotated with
+    `handles_ambiguity`. In this case, the function will receive a list
+    of all the transformations of all the derivations of the rule.
+    By default, a list of trees where each tree.data is equal to the
+    rule name or one of its aliases.
+
+    Non-tree transformations are made possible by override of
+    `__default__`, `__default_token__`, and `__default_ambig__`.
+    """
+
+    def __init__(self, tree_class=Tree, prioritizer=ForestSumVisitor(), resolve_ambiguity=True):
+        super(TreeForestTransformer, self).__init__(tree_class, dict(), prioritizer, resolve_ambiguity)
+
+    def __default__(self, name, data):
+        """Default operation on tree (for override).
+
+        Returns a tree with name with data as children.
+        """
+        return self.tree_class(name, data)
+
+    def __default_ambig__(self, name, data):
+        """Default operation on ambiguous rule (for override).
+
+        Wraps data in an '_ambig_ node if it contains more than
+        one element.'
+        """
+        if len(data) > 1:
+            return self.tree_class('_ambig', data)
+        elif data:
+            return data[0]
+        raise Discard
+
+    def __default_token__(self, node):
+        """Default operation on Token (for override).
+
+        Returns node
+        """
+        return node
+
+    def transform_token_node(self, node):
+        return getattr(self, node.type, self.__default_token__)(node)
+
+    def _call_rule_func(self, node, data):
+        name = node.rule.alias or node.rule.options.template_source or node.rule.origin.name
+        user_func = getattr(self, name, self.__default__)
+        if user_func == self.__default__ or hasattr(user_func, 'handles_ambiguity'):
+            user_func = partial(self.__default__, name)
+        if not self.resolve_ambiguity:
+            wrapper = partial(AmbiguousIntermediateExpander, self.tree_class)
+            user_func = wrapper(user_func)
+        return user_func(data)
+
+    def _call_ambig_func(self, node, data):
+        name = node.s.name
+        user_func = getattr(self, name, self.__default_ambig__)
+        if user_func == self.__default_ambig__ or not hasattr(user_func, 'handles_ambiguity'):
+            user_func = partial(self.__default_ambig__, name)
+        return user_func(data)
+
 class ForestToPyDotVisitor(ForestVisitor):
     """
     A Forest visitor which writes the SPPF to a PNG.