1"""
2    pint.pint_eval
3    ~~~~~~~~~~~~~~
4
5    An expression evaluator to be used as a safe replacement for builtin eval.
6
7    :copyright: 2016 by Pint Authors, see AUTHORS for more details.
8    :license: BSD, see LICENSE for more details.
9"""
10
11import operator
12import token as tokenlib
13
14from .errors import DefinitionSyntaxError
15
16# For controlling order of operations
17_OP_PRIORITY = {
18    "**": 3,
19    "^": 3,
20    "unary": 2,
21    "*": 1,
22    "": 1,  # operator for implicit ops
23    "/": 1,
24    "+": 0,
25    "-": 0,
26}
27
28
29def _power(left, right):
30    from .compat import is_duck_array
31    from .quantity import Quantity
32
33    if (
34        isinstance(left, Quantity)
35        and is_duck_array(left.magnitude)
36        and left.dtype.kind not in "cf"
37        and right < 0
38    ):
39        left = left.astype(float)
40
41    return operator.pow(left, right)
42
43
44_BINARY_OPERATOR_MAP = {
45    "**": _power,
46    "*": operator.mul,
47    "": operator.mul,  # operator for implicit ops
48    "/": operator.truediv,
49    "+": operator.add,
50    "-": operator.sub,
51}
52
53_UNARY_OPERATOR_MAP = {"+": lambda x: x, "-": lambda x: x * -1}
54
55
56class EvalTreeNode:
57    """Single node within an evaluation tree
58
59    left + operator + right --> binary op
60    left + operator --> unary op
61    left + right --> implicit op
62    left --> single value
63    """
64
65    def __init__(self, left, operator=None, right=None):
66        self.left = left
67        self.operator = operator
68        self.right = right
69
70    def to_string(self):
71        # For debugging purposes
72        if self.right:
73            comps = [self.left.to_string()]
74            if self.operator:
75                comps.append(self.operator[1])
76            comps.append(self.right.to_string())
77        elif self.operator:
78            comps = [self.operator[1], self.left.to_string()]
79        else:
80            return self.left[1]
81        return "(%s)" % " ".join(comps)
82
83    def evaluate(self, define_op, bin_op=None, un_op=None):
84        """Evaluate node.
85
86        Parameters
87        ----------
88        define_op : callable
89            Translates tokens into objects.
90        bin_op : dict or None, optional
91             (Default value = _BINARY_OPERATOR_MAP)
92        un_op : dict or None, optional
93             (Default value = _UNARY_OPERATOR_MAP)
94
95        Returns
96        -------
97
98        """
99
100        bin_op = bin_op or _BINARY_OPERATOR_MAP
101        un_op = un_op or _UNARY_OPERATOR_MAP
102
103        if self.right:
104            # binary or implicit operator
105            op_text = self.operator[1] if self.operator else ""
106            if op_text not in bin_op:
107                raise DefinitionSyntaxError('missing binary operator "%s"' % op_text)
108            left = self.left.evaluate(define_op, bin_op, un_op)
109            return bin_op[op_text](left, self.right.evaluate(define_op, bin_op, un_op))
110        elif self.operator:
111            # unary operator
112            op_text = self.operator[1]
113            if op_text not in un_op:
114                raise DefinitionSyntaxError('missing unary operator "%s"' % op_text)
115            return un_op[op_text](self.left.evaluate(define_op, bin_op, un_op))
116        else:
117            # single value
118            return define_op(self.left)
119
120
121def build_eval_tree(tokens, op_priority=_OP_PRIORITY, index=0, depth=0, prev_op=None):
122    """Build an evaluation tree from a set of tokens.
123
124    Params:
125    Index, depth, and prev_op used recursively, so don't touch.
126    Tokens is an iterable of tokens from an expression to be evaluated.
127
128    Transform the tokens from an expression into a recursive parse tree, following order
129    of operations. Operations can include binary ops (3 + 4), implicit ops (3 kg), or
130    unary ops (-1).
131
132    General Strategy:
133    1) Get left side of operator
134    2) If no tokens left, return final result
135    3) Get operator
136    4) Use recursion to create tree starting at token on right side of operator (start at step #1)
137    4.1) If recursive call encounters an operator with lower or equal priority to step #2, exit recursion
138    5) Combine left side, operator, and right side into a new left side
139    6) Go back to step #2
140
141    """
142
143    if depth == 0 and prev_op is None:
144        # ensure tokens is list so we can access by index
145        tokens = list(tokens)
146
147    result = None
148
149    while True:
150        current_token = tokens[index]
151        token_type = current_token[0]
152        token_text = current_token[1]
153
154        if token_type == tokenlib.OP:
155            if token_text == ")":
156                if prev_op is None:
157                    raise DefinitionSyntaxError(
158                        "unopened parentheses in tokens: %s" % current_token
159                    )
160                elif prev_op == "(":
161                    # close parenthetical group
162                    return result, index
163                else:
164                    # parenthetical group ending, but we need to close sub-operations within group
165                    return result, index - 1
166            elif token_text == "(":
167                # gather parenthetical group
168                right, index = build_eval_tree(
169                    tokens, op_priority, index + 1, 0, token_text
170                )
171                if not tokens[index][1] == ")":
172                    raise DefinitionSyntaxError("weird exit from parentheses")
173                if result:
174                    # implicit op with a parenthetical group, i.e. "3 (kg ** 2)"
175                    result = EvalTreeNode(left=result, right=right)
176                else:
177                    # get first token
178                    result = right
179            elif token_text in op_priority:
180                if result:
181                    # equal-priority operators are grouped in a left-to-right order,
182                    # unless they're exponentiation, in which case they're grouped
183                    # right-to-left this allows us to get the expected behavior for
184                    # multiple exponents
185                    #     (2^3^4)  --> (2^(3^4))
186                    #     (2 * 3 / 4) --> ((2 * 3) / 4)
187                    if op_priority[token_text] <= op_priority.get(
188                        prev_op, -1
189                    ) and token_text not in ["**", "^"]:
190                        # previous operator is higher priority, so end previous binary op
191                        return result, index - 1
192                    # get right side of binary op
193                    right, index = build_eval_tree(
194                        tokens, op_priority, index + 1, depth + 1, token_text
195                    )
196                    result = EvalTreeNode(
197                        left=result, operator=current_token, right=right
198                    )
199                else:
200                    # unary operator
201                    right, index = build_eval_tree(
202                        tokens, op_priority, index + 1, depth + 1, "unary"
203                    )
204                    result = EvalTreeNode(left=right, operator=current_token)
205        elif token_type == tokenlib.NUMBER or token_type == tokenlib.NAME:
206            if result:
207                # tokens with an implicit operation i.e. "1 kg"
208                if op_priority[""] <= op_priority.get(prev_op, -1):
209                    # previous operator is higher priority than implicit, so end
210                    # previous binary op
211                    return result, index - 1
212                right, index = build_eval_tree(
213                    tokens, op_priority, index, depth + 1, ""
214                )
215                result = EvalTreeNode(left=result, right=right)
216            else:
217                # get first token
218                result = EvalTreeNode(left=current_token)
219
220        if tokens[index][0] == tokenlib.ENDMARKER:
221            if prev_op == "(":
222                raise DefinitionSyntaxError("unclosed parentheses in tokens")
223            if depth > 0 or prev_op:
224                # have to close recursion
225                return result, index
226            else:
227                # recursion all closed, so just return the final result
228                return result
229
230        if index + 1 >= len(tokens):
231            # should hit ENDMARKER before this ever happens
232            raise DefinitionSyntaxError("unexpected end to tokens")
233
234        index += 1
235