| 10 | position depends on previous tokens and prediction of newline. |
| 11 | """ |
| 12 | class ProcessTokens(ParseTreeListener): |
| 13 | def __init__(self, newline_forest, indent_forest, whitespace_forest, vec, stream): |
| 14 | self.newline_forest = newline_forest |
| 15 | self.indent_forest = indent_forest |
| 16 | self.whitespace_forest = whitespace_forest |
| 17 | self.vec = vec # the DictVectorizer used to transform categorical features |
| 18 | self.stream = stream # track stream so we can examine previous tokens |
| 19 | self.current_indent = 0 |
| 20 | |
| 21 | def visitTerminal(self, node): |
| 22 | curToken = node.symbol |
| 23 | if curToken.type==-1: |
| 24 | return |
| 25 | if curToken.tokenIndex>=1: |
| 26 | # predict newline based upon curToken appearing after prevToken on same line |
| 27 | prevToken = self.stream.tokens[curToken.tokenIndex-1] |
| 28 | curLine = prevToken.line |
| 29 | curToken.column = prevToken.column + len(prevToken.text) |
| 30 | else: |
| 31 | curLine = 1 |
| 32 | curToken.column = 0 |
| 33 | curToken.line = curLine |
| 34 | |
| 35 | vars = groomlib.node_features(self.stream.tokens, node) |
| 36 | |
| 37 | feature_names = groomlib.FEATURE_NAMES |
| 38 | n = len(feature_names) |
| 39 | d = dict((feature_names[i], vars[i]) for i in range(0, n)) |
| 40 | # print d |
| 41 | transformed_features_testing = self.vec.transform(d).toarray() |
| 42 | inject_newline = self.newline_forest.predict(transformed_features_testing) |
| 43 | # newline_predictions_proba = self.forest.predict_proba(transformed_data_testing) |
| 44 | # print curToken, "->", newline_predictions_proba |
| 45 | # print "inject_newline", inject_newline |
| 46 | if inject_newline: |
| 47 | curToken.line += 1 |
| 48 | curToken.column = 0 |
| 49 | print # inject a newline |
| 50 | # now figure out indent |
| 51 | indent = self.indent_forest.predict(transformed_features_testing) |
| 52 | # if indent!=0: |
| 53 | # print "indent %d at %s" % (indent,curToken) |
| 54 | self.current_indent += indent |
| 55 | sys.stdout.write(" " * self.current_indent) # inject indent |
| 56 | else: |
| 57 | ws = self.whitespace_forest.predict(transformed_features_testing) |
| 58 | sys.stdout.write(" " * ws) # inject whitespace before token |
| 59 | |
| 60 | sys.stdout.write(curToken.text) |