1from functools import wraps
2
3from .utils import smart_decorator, combine_alternatives
4from .tree import Tree
5from .exceptions import VisitError, GrammarError
6from .lexer import Token
7
8###{standalone
9from inspect import getmembers, getmro
10
11class Discard(Exception):
12    pass
13
14# Transformers
15
16class _Decoratable:
17    @classmethod
18    def _apply_decorator(cls, decorator, **kwargs):
19        mro = getmro(cls)
20        assert mro[0] is cls
21        libmembers = {name for _cls in mro[1:] for name, _ in getmembers(_cls)}
22        for name, value in getmembers(cls):
23
24            # Make sure the function isn't inherited (unless it's overwritten)
25            if name.startswith('_') or (name in libmembers and name not in cls.__dict__):
26                continue
27            if not callable(value):
28                continue
29
30            # Skip if v_args already applied (at the function level)
31            if hasattr(cls.__dict__[name], 'vargs_applied') or hasattr(value, 'vargs_applied'):
32                continue
33
34            static = isinstance(cls.__dict__[name], (staticmethod, classmethod))
35            setattr(cls, name, decorator(value, static=static, **kwargs))
36        return cls
37
38    def __class_getitem__(cls, _):
39        return cls
40
41
42class Transformer(_Decoratable):
43    """Visits the tree recursively, starting with the leaves and finally the root (bottom-up)
44
45    Calls its methods (provided by user via inheritance) according to tree.data
46    The returned value replaces the old one in the structure.
47
48    Can be used to implement map or reduce.
49    """
50    __visit_tokens__ = True   # For backwards compatibility
51
52    def __init__(self,  visit_tokens=True):
53        self.__visit_tokens__ = visit_tokens
54
55    def _call_userfunc(self, tree, new_children=None):
56        # Assumes tree is already transformed
57        children = new_children if new_children is not None else tree.children
58        try:
59            f = getattr(self, tree.data)
60        except AttributeError:
61            return self.__default__(tree.data, children, tree.meta)
62        else:
63            try:
64                wrapper = getattr(f, 'visit_wrapper', None)
65                if wrapper is not None:
66                    return f.visit_wrapper(f, tree.data, children, tree.meta)
67                else:
68                    return f(children)
69            except (GrammarError, Discard):
70                raise
71            except Exception as e:
72                raise VisitError(tree.data, tree, e)
73
74    def _call_userfunc_token(self, token):
75        try:
76            f = getattr(self, token.type)
77        except AttributeError:
78            return self.__default_token__(token)
79        else:
80            try:
81                return f(token)
82            except (GrammarError, Discard):
83                raise
84            except Exception as e:
85                raise VisitError(token.type, token, e)
86
87
88    def _transform_children(self, children):
89        for c in children:
90            try:
91                if isinstance(c, Tree):
92                    yield self._transform_tree(c)
93                elif self.__visit_tokens__ and isinstance(c, Token):
94                    yield self._call_userfunc_token(c)
95                else:
96                    yield c
97            except Discard:
98                pass
99
100    def _transform_tree(self, tree):
101        children = list(self._transform_children(tree.children))
102        return self._call_userfunc(tree, children)
103
104    def transform(self, tree):
105        return self._transform_tree(tree)
106
107    def __mul__(self, other):
108        return TransformerChain(self, other)
109
110    def __default__(self, data, children, meta):
111        "Default operation on tree (for override)"
112        return Tree(data, children, meta)
113
114    def __default_token__(self, token):
115        "Default operation on token (for override)"
116        return token
117
118
119
120class InlineTransformer(Transformer):   # XXX Deprecated
121    def _call_userfunc(self, tree, new_children=None):
122        # Assumes tree is already transformed
123        children = new_children if new_children is not None else tree.children
124        try:
125            f = getattr(self, tree.data)
126        except AttributeError:
127            return self.__default__(tree.data, children, tree.meta)
128        else:
129            return f(*children)
130
131
132class TransformerChain(object):
133    def __init__(self, *transformers):
134        self.transformers = transformers
135
136    def transform(self, tree):
137        for t in self.transformers:
138            tree = t.transform(tree)
139        return tree
140
141    def __mul__(self, other):
142        return TransformerChain(*self.transformers + (other,))
143
144
145class Transformer_InPlace(Transformer):
146    "Non-recursive. Changes the tree in-place instead of returning new instances"
147    def _transform_tree(self, tree):           # Cancel recursion
148        return self._call_userfunc(tree)
149
150    def transform(self, tree):
151        for subtree in tree.iter_subtrees():
152            subtree.children = list(self._transform_children(subtree.children))
153
154        return self._transform_tree(tree)
155
156
157class Transformer_NonRecursive(Transformer):
158    "Non-recursive. Doesn't change the original tree."
159
160    def transform(self, tree):
161        # Tree to postfix
162        rev_postfix = []
163        q = [tree]
164        while q:
165            t = q.pop()
166            rev_postfix.append( t )
167            if isinstance(t, Tree):
168                q += t.children
169
170        # Postfix to tree
171        stack = []
172        for x in reversed(rev_postfix):
173            if isinstance(x, Tree):
174                size = len(x.children)
175                if size:
176                    args = stack[-size:]
177                    del stack[-size:]
178                else:
179                    args = []
180                stack.append(self._call_userfunc(x, args))
181            else:
182                stack.append(x)
183
184        t ,= stack  # We should have only one tree remaining
185        return t
186
187
188
189class Transformer_InPlaceRecursive(Transformer):
190    "Recursive. Changes the tree in-place instead of returning new instances"
191    def _transform_tree(self, tree):
192        tree.children = list(self._transform_children(tree.children))
193        return self._call_userfunc(tree)
194
195
196
197# Visitors
198
199class VisitorBase:
200    def _call_userfunc(self, tree):
201        return getattr(self, tree.data, self.__default__)(tree)
202
203    def __default__(self, tree):
204        "Default operation on tree (for override)"
205        return tree
206
207    def __class_getitem__(cls, _):
208        return cls
209
210
211class Visitor(VisitorBase):
212    """Bottom-up visitor, non-recursive
213
214    Visits the tree, starting with the leaves and finally the root (bottom-up)
215    Calls its methods (provided by user via inheritance) according to tree.data
216    """
217
218    def visit(self, tree):
219        for subtree in tree.iter_subtrees():
220            self._call_userfunc(subtree)
221        return tree
222
223    def visit_topdown(self,tree):
224        for subtree in tree.iter_subtrees_topdown():
225            self._call_userfunc(subtree)
226        return tree
227
228class Visitor_Recursive(VisitorBase):
229    """Bottom-up visitor, recursive
230
231    Visits the tree, starting with the leaves and finally the root (bottom-up)
232    Calls its methods (provided by user via inheritance) according to tree.data
233    """
234
235    def visit(self, tree):
236        for child in tree.children:
237            if isinstance(child, Tree):
238                self.visit(child)
239
240        self._call_userfunc(tree)
241        return tree
242
243    def visit_topdown(self,tree):
244        self._call_userfunc(tree)
245
246        for child in tree.children:
247            if isinstance(child, Tree):
248                self.visit_topdown(child)
249
250        return tree
251
252
253
254def visit_children_decor(func):
255    "See Interpreter"
256    @wraps(func)
257    def inner(cls, tree):
258        values = cls.visit_children(tree)
259        return func(cls, values)
260    return inner
261
262
263class Interpreter(_Decoratable):
264    """Top-down visitor, recursive
265
266    Visits the tree, starting with the root and finally the leaves (top-down)
267    Calls its methods (provided by user via inheritance) according to tree.data
268
269    Unlike Transformer and Visitor, the Interpreter doesn't automatically visit its sub-branches.
270    The user has to explicitly call visit_children, or use the @visit_children_decor
271    """
272
273    def visit(self, tree):
274        f = getattr(self, tree.data)
275        wrapper = getattr(f, 'visit_wrapper', None)
276        if wrapper is not None:
277            return f.visit_wrapper(f, tree.data, tree.children, tree.meta)
278        else:
279            return f(tree)
280
281    def visit_children(self, tree):
282        return [self.visit(child) if isinstance(child, Tree) else child
283                for child in tree.children]
284
285    def __getattr__(self, name):
286        return self.__default__
287
288    def __default__(self, tree):
289        return self.visit_children(tree)
290
291
292
293
294# Decorators
295
296def _apply_decorator(obj, decorator, **kwargs):
297    try:
298        _apply = obj._apply_decorator
299    except AttributeError:
300        return decorator(obj, **kwargs)
301    else:
302        return _apply(decorator, **kwargs)
303
304
305
306def _inline_args__func(func):
307    @wraps(func)
308    def create_decorator(_f, with_self):
309        if with_self:
310            def f(self, children):
311                return _f(self, *children)
312        else:
313            def f(self, children):
314                return _f(*children)
315        return f
316
317    return smart_decorator(func, create_decorator)
318
319
320def inline_args(obj):   # XXX Deprecated
321    return _apply_decorator(obj, _inline_args__func)
322
323
324
325def _visitor_args_func_dec(func, visit_wrapper=None, static=False):
326    def create_decorator(_f, with_self):
327        if with_self:
328            def f(self, *args, **kwargs):
329                return _f(self, *args, **kwargs)
330        else:
331            def f(self, *args, **kwargs):
332                return _f(*args, **kwargs)
333        return f
334
335    if static:
336        f = wraps(func)(create_decorator(func, False))
337    else:
338        f = smart_decorator(func, create_decorator)
339    f.vargs_applied = True
340    f.visit_wrapper = visit_wrapper
341    return f
342
343
344def _vargs_inline(f, data, children, meta):
345    return f(*children)
346def _vargs_meta_inline(f, data, children, meta):
347    return f(meta, *children)
348def _vargs_meta(f, data, children, meta):
349    return f(children, meta)   # TODO swap these for consistency? Backwards incompatible!
350def _vargs_tree(f, data, children, meta):
351    return f(Tree(data, children, meta))
352
353def v_args(inline=False, meta=False, tree=False, wrapper=None):
354    "A convenience decorator factory, for modifying the behavior of user-supplied visitor methods"
355    if tree and (meta or inline):
356        raise ValueError("Visitor functions cannot combine 'tree' with 'meta' or 'inline'.")
357
358    func = None
359    if meta:
360        if inline:
361            func = _vargs_meta_inline
362        else:
363            func = _vargs_meta
364    elif inline:
365        func = _vargs_inline
366    elif tree:
367        func = _vargs_tree
368
369    if wrapper is not None:
370        if func is not None:
371            raise ValueError("Cannot use 'wrapper' along with 'tree', 'meta' or 'inline'.")
372        func = wrapper
373
374    def _visitor_args_dec(obj):
375        return _apply_decorator(obj, _visitor_args_func_dec, visit_wrapper=func)
376    return _visitor_args_dec
377
378
379###}
380
381
382#--- Visitor Utilities ---
383
384class CollapseAmbiguities(Transformer):
385    """
386    Transforms a tree that contains any number of _ambig nodes into a list of trees,
387    each one containing an unambiguous tree.
388
389    The length of the resulting list is the product of the length of all _ambig nodes.
390
391    Warning: This may quickly explode for highly ambiguous trees.
392
393    """
394    def _ambig(self, options):
395        return sum(options, [])
396    def __default__(self, data, children_lists, meta):
397        return [Tree(data, children, meta) for children in combine_alternatives(children_lists)]
398    def __default_token__(self, t):
399        return [t]
400