forked from pydata/patsy
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathparse_formula.py
More file actions
298 lines (253 loc) · 9.52 KB
/
parse_formula.py
File metadata and controls
298 lines (253 loc) · 9.52 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
# This file is part of Patsy
# Copyright (C) 2011 Nathaniel Smith <njs@pobox.com>
# See file LICENSE.txt for license information.
# This file defines a parser for a simple language based on S/R "formulas"
# (which are described in sections 2.3 and 2.4 in Chambers & Hastie, 1992). It
# uses the machinery in patsy.parse_core to do the heavy-lifting -- its
# biggest job is to handle tokenization.
from __future__ import print_function
__all__ = ["parse_formula"]
# The Python tokenizer
import tokenize
import six
from six.moves import cStringIO as StringIO
from patsy import PatsyError
from patsy.origin import Origin
from patsy.infix_parser import Token, Operator, infix_parse, ParseNode
from patsy.tokens import python_tokenize, pretty_untokenize
from patsy.util import PushbackAdapter
_atomic_token_types = ["PYTHON_EXPR", "ZERO", "ONE", "NUMBER"]
def _is_a(f, v):
try:
f(v)
except ValueError:
return False
else:
return True
# Helper function for _tokenize_formula:
def _read_python_expr(it, end_tokens):
# Read out a full python expression, stopping when we hit an
# unnested end token.
pytypes = []
token_strings = []
origins = []
bracket_level = 0
for pytype, token_string, origin in it:
assert bracket_level >= 0
if bracket_level == 0 and token_string in end_tokens:
it.push_back((pytype, token_string, origin))
break
if token_string in ("(", "[", "{"):
bracket_level += 1
if token_string in (")", "]", "}"):
bracket_level -= 1
if bracket_level < 0:
raise PatsyError("unmatched close bracket", origin)
pytypes.append(pytype)
token_strings.append(token_string)
origins.append(origin)
# Either we found an end_token, or we hit the end of the string
if bracket_level == 0:
expr_text = pretty_untokenize(zip(pytypes, token_strings))
if expr_text == "0":
token_type = "ZERO"
elif expr_text == "1":
token_type = "ONE"
elif _is_a(int, expr_text) or _is_a(float, expr_text):
token_type = "NUMBER"
else:
token_type = "PYTHON_EXPR"
return Token(token_type, Origin.combine(origins), extra=expr_text)
else:
raise PatsyError("unclosed bracket in embedded Python "
"expression",
Origin.combine(origins))
def _tokenize_formula(code, operator_strings):
assert "(" not in operator_strings
assert ")" not in operator_strings
magic_token_types = {"(": Token.LPAREN,
")": Token.RPAREN,
}
for operator_string in operator_strings:
magic_token_types[operator_string] = operator_string
# Once we enter a Python expression, a ( does not end it, but any other
# "magic" token does:
end_tokens = set(magic_token_types)
end_tokens.remove("(")
it = PushbackAdapter(python_tokenize(code))
for pytype, token_string, origin in it:
if token_string in magic_token_types:
yield Token(magic_token_types[token_string], origin)
else:
it.push_back((pytype, token_string, origin))
yield _read_python_expr(it, end_tokens)
def test__tokenize_formula():
code = "y ~ a + (foo(b,c + 2)) + -1 + 0 + 10"
tokens = list(_tokenize_formula(code, ["+", "-", "~"]))
expecteds = [("PYTHON_EXPR", Origin(code, 0, 1), "y"),
("~", Origin(code, 2, 3), None),
("PYTHON_EXPR", Origin(code, 4, 5), "a"),
("+", Origin(code, 6, 7), None),
(Token.LPAREN, Origin(code, 8, 9), None),
("PYTHON_EXPR", Origin(code, 9, 23), "foo(b, c + 2)"),
(Token.RPAREN, Origin(code, 23, 24), None),
("+", Origin(code, 25, 26), None),
("-", Origin(code, 27, 28), None),
("ONE", Origin(code, 28, 29), "1"),
("+", Origin(code, 30, 31), None),
("ZERO", Origin(code, 32, 33), "0"),
("+", Origin(code, 34, 35), None),
("NUMBER", Origin(code, 36, 38), "10"),
]
for got, expected in zip(tokens, expecteds):
assert isinstance(got, Token)
assert got.type == expected[0]
assert got.origin == expected[1]
assert got.extra == expected[2]
_unary_tilde = Operator("~", 1, -100)
_default_ops = [
_unary_tilde,
Operator("~", 2, -100),
Operator("+", 2, 100),
Operator("-", 2, 100),
Operator("*", 2, 200),
Operator("/", 2, 200),
Operator(":", 2, 300),
Operator("**", 2, 500),
Operator("+", 1, 100),
Operator("-", 1, 100),
]
def parse_formula(code, extra_operators=[]):
if not code.strip():
code = "~ 1"
for op in extra_operators:
if op.precedence < 0:
raise ValueError("all operators must have precedence >= 0")
operators = _default_ops + extra_operators
operator_strings = [op.token_type for op in operators]
tree = infix_parse(_tokenize_formula(code, operator_strings),
operators,
_atomic_token_types)
if not isinstance(tree, ParseNode) or tree.type != "~":
tree = ParseNode("~", None, [tree], tree.origin)
return tree
#############
_parser_tests = {
"": ["~", "1"],
" ": ["~", "1"],
" \n ": ["~", "1"],
"1": ["~", "1"],
"a": ["~", "a"],
"a ~ b": ["~", "a", "b"],
"(a ~ b)": ["~", "a", "b"],
"a ~ ((((b))))": ["~", "a", "b"],
"a ~ ((((+b))))": ["~", "a", ["+", "b"]],
"a + b + c": ["~", ["+", ["+", "a", "b"], "c"]],
"a + (b ~ c) + d": ["~", ["+", ["+", "a", ["~", "b", "c"]], "d"]],
"a + np.log(a, base=10)": ["~", ["+", "a", "np.log(a, base=10)"]],
# Note different spacing:
"a + np . log(a , base = 10)": ["~", ["+", "a", "np.log(a, base=10)"]],
# Check precedence
"a + b ~ c * d": ["~", ["+", "a", "b"], ["*", "c", "d"]],
"a + b * c": ["~", ["+", "a", ["*", "b", "c"]]],
"-a**2": ["~", ["-", ["**", "a", "2"]]],
"-a:b": ["~", ["-", [":", "a", "b"]]],
"a + b:c": ["~", ["+", "a", [":", "b", "c"]]],
"(a + b):c": ["~", [":", ["+", "a", "b"], "c"]],
"a*b:c": ["~", ["*", "a", [":", "b", "c"]]],
"a+b / c": ["~", ["+", "a", ["/", "b", "c"]]],
"~ a": ["~", "a"],
"-1": ["~", ["-", "1"]],
}
def _compare_trees(got, expected):
assert isinstance(got, ParseNode)
if got.args:
assert got.type == expected[0]
for arg, expected_arg in zip(got.args, expected[1:]):
_compare_trees(arg, expected_arg)
else:
assert got.type in _atomic_token_types
assert got.token.extra == expected
def _do_parse_test(test_cases, extra_operators):
for code, expected in six.iteritems(test_cases):
actual = parse_formula(code, extra_operators=extra_operators)
print(repr(code), repr(expected))
print(actual)
_compare_trees(actual, expected)
def test_parse_formula():
_do_parse_test(_parser_tests, [])
def test_parse_origin():
tree = parse_formula("a ~ b + c")
assert tree.origin == Origin("a ~ b + c", 0, 9)
assert tree.token.origin == Origin("a ~ b + c", 2, 3)
assert tree.args[0].origin == Origin("a ~ b + c", 0, 1)
assert tree.args[1].origin == Origin("a ~ b + c", 4, 9)
assert tree.args[1].token.origin == Origin("a ~ b + c", 6, 7)
assert tree.args[1].args[0].origin == Origin("a ~ b + c", 4, 5)
assert tree.args[1].args[1].origin == Origin("a ~ b + c", 8, 9)
# <> mark off where the error should be reported:
_parser_error_tests = [
"a <+>",
"a + <(>",
"a + b <# asdf>",
"<)>",
"a + <)>",
"<*> a",
"a + <*>",
"a + <foo[bar>",
"a + <foo{bar>",
"a + <foo(bar>",
"a + <[bar>",
"a + <{bar>",
"a + <{bar[]>",
"a + foo<]>bar",
"a + foo[]<]>bar",
"a + foo{}<}>bar",
"a + foo<)>bar",
"a + b<)>",
"(a) <.>",
"<(>a + b",
"a +< >'foo", # Not the best placement for the error
]
# Split out so it can also be used by tests of the evaluator (which also
# raises PatsyError's)
def _parsing_error_test(parse_fn, error_descs): # pragma: no cover
for error_desc in error_descs:
letters = []
start = None
end = None
for letter in error_desc:
if letter == "<":
start = len(letters)
elif letter == ">":
end = len(letters)
else:
letters.append(letter)
bad_code = "".join(letters)
assert start is not None and end is not None
print(error_desc)
print(repr(bad_code), start, end)
try:
parse_fn(bad_code)
except PatsyError as e:
print(e)
assert e.origin.code == bad_code
assert e.origin.start == start
assert e.origin.end == end
else:
assert False, "parser failed to report an error!"
def test_parse_errors(extra_operators=[]):
def parse_fn(code):
return parse_formula(code, extra_operators=extra_operators)
_parsing_error_test(parse_fn, _parser_error_tests)
_extra_op_parser_tests = {
"a | b": ["~", ["|", "a", "b"]],
"a * b|c": ["~", ["*", "a", ["|", "b", "c"]]],
}
def test_parse_extra_op():
extra_operators = [Operator("|", 2, 250)]
_do_parse_test(_parser_tests,
extra_operators=extra_operators)
_do_parse_test(_extra_op_parser_tests,
extra_operators=extra_operators)
test_parse_errors(extra_operators=extra_operators)