| 2 | from antlr4 import ParseTreeListener |
| 3 | |
| 4 | class CollectTokenFeatures(ParseTreeListener): |
| 5 | def __init__(self, stream): |
| 6 | self.stream = stream # track stream so we can examine previous tokens |
| 7 | self.inject_newlines = [] # prediction |
| 8 | self.indent = [] # prediction of indent/dedent == column delta |
| 9 | self.whitespace = [] # prediction of whitespace before token |
| 10 | self.features = [] # independent vars |
| 11 | self.first_token_on_line = None # track to compute indent |
| 12 | |
| 13 | def visitTerminal(self, node): |
| 14 | i = node.symbol.tokenIndex |
| 15 | curToken = node.symbol |
| 16 | if curToken.type==-1: |
| 17 | return |
| 18 | precedingNL = False |
| 19 | column_delta = 0 |
| 20 | ws = 0 |
| 21 | if i>=1: |
| 22 | prevToken = self.stream.tokens[i-1] |
| 23 | precedingNL = curToken.line > prevToken.line |
| 24 | if precedingNL: |
| 25 | if self.first_token_on_line is not None: |
| 26 | column_delta = curToken.column - self.first_token_on_line.column |
| 27 | self.first_token_on_line = curToken |
| 28 | else: |
| 29 | ws = curToken.column - (prevToken.column+len(prevToken.text)) |
| 30 | |
| 31 | #print dir(node) |
| 32 | vars = groomlib.node_features(self.stream.tokens, node) |
| 33 | self.inject_newlines.append(1 if precedingNL else 0) |
| 34 | self.indent.append(column_delta) |
| 35 | self.whitespace.append(ws) |
| 36 | self.features.append(vars) |
no outgoing calls
no test coverage detected