from io import StringIO from atheris.token import TokenType, Token class Scanner(object): def __init__(self, obj): buf = None if isinstance(obj, str): buf = StringIO(obj) else: buf = obj self.buf = buf def readline(self): return self.buf.readline() def read(self): return self.buf.read(1) def unread(self): self.buf.seek(self.buf.tell()-1) class Lexer(object): def __init__(self, obj): self.scanner = Scanner(obj) self.pos = 0 self.col = 0 self.line = 0 self.line_buf = '' self.cur_char = '' self.indent_stack = [] self.advance() def fill_buffer(self): self.line_buf = self.scanner.readline().rstrip('\n').rstrip('\r\n') def advance(self, return_on_newline=False): self.pos += 1 self.col += 1 end_of_line = False if self.pos < len(self.line_buf): self.cur_char = self.line_buf[self.pos] else: end_of_line = True if return_on_newline: return end_of_line # if we've reached the end of the line, fill er up again self.fill_buffer() if self.line_buf: self.pos = 0 self.line += 1 self.col = 1 self.cur_char = self.line_buf[self.pos] else: # no more characters self.cur_char = '' return end_of_line def get_indent(self): indent = 0 end_of_line = False while self.cur_char.isspace() and self.cur_char and not end_of_line: indent += 1 end_of_line = self.advance(True) token = None if self.indent_stack: last_indent = self.indent_stack[-1] if indent < last_indent: self.indent_stack.pop() size = last_indent - indent self.col -= 1 token = Token(TokenType.DEDENT, size, self.line, self.col) self.advance() elif indent > last_indent: self.indent_stack.append(indent) size = indent - last_indent self.col -= 1 token = Token(TokenType.INDENT, size, self.line, self.col-size) self.advance() elif indent > 0: self.indent_stack.append(indent) self.col -= 1 token = Token(TokenType.INDENT, indent, self.line, 1) self.advance() return token def next_token(self): cur_col = self.col cur_line = self.line if self.cur_char: if self.pos == 0: indent = self.get_indent() if indent is not None: return indent while self.cur_char.isspace(): self.advance() if self.cur_char.isdigit() or self.cur_char == '.': num_str = '' while self.cur_char.isdigit() or self.cur_char == '.': num_str += self.cur_char self.advance() return Token(TokenType.NUMBER, num_str, cur_line, cur_col) elif self.cur_char == '#': self.advance() while self.cur_char and self.cur_char not in set(['\n', '\r']): self.advance() return self.next_token() elif self.cur_char in Token.keyword_map: last = self.cur_char self.advance() return Token(Token.keyword_map[last], last, cur_line, cur_col) elif not self.cur_char.isspace(): id_str = '' while not self.cur_char.isspace() and not self.cur_char in Token.keyword_map: id_str += self.cur_char self.advance() if self.cur_char in Token.keyword_map and not id_str: token = Token(Token.keyword_map[self.cur_char], self.cur_char, cur_line, cur_col) self.advance() return token if id_str in Token.keyword_map: return Token(Token.keyword_map[id_str], id_str, cur_line, cur_col) else: return Token(TokenType.IDENT, id_str, cur_line, cur_col) if self.indent_stack: # Extra indents that have not been taken care of (ie: at the # end of an indented file) indent = self.indent_stack.pop() content = indent if self.indent_stack: previous_indent = self.indent_stack[-1] content = indent - previous_indent return Token(TokenType.DEDENT, content, cur_line+1, 1) return Token(TokenType.EOF, '', cur_line+1, 1) def tokens(self): token = self.next_token() while token.token_type != TokenType.EOF: yield token token = self.next_token() yield token if __name__ == '__main__': import sys lexer = Lexer(open(sys.argv[1])) for token in lexer.tokens(): print(token)