Simple Arithmetic Lexer

lexer.py

import collections
import enum
import re

from sre_parse import Pattern, SubPattern, parse
from sre_compile import compile as sre_compile
from sre_constants import BRANCH, SUBPATTERN


class Tokens(enum.Enum):

    WHITESPACE = 1
    PLUS = 2
    MINUS = 3
    MUL = 4
    DIV = 5
    NUM = 6
    LPAREN = 7
    RPAREN = 8


Token = collections.namedtuple('Token', ('type', 'value'))



SCANNER = re.Scanner([
    (r'\+', lambda s, token: Token(Tokens.PLUS, token)),
    ('-', lambda s, token: Token(Tokens.MINUS, token)),
    (r'\*', lambda s, token: Token(Tokens.MUL, token)),
    ('/', lambda s, token: Token(Tokens.DIV, token)),
    (r'\d+', lambda s, token: Token(Tokens.NUM, token)),
    (r'\(', lambda s, token: Token(Tokens.LPAREN, token)),
    (r'\)', lambda s, token: Token(Tokens.RPAREN, token)),
    (r'\s+', None),
])


tokenize = SCANNER.scan


def test_tokenize():
    result = list(tokenize('(1 + 2) * 3')[0])
    expected = [
        Token(Tokens.LPAREN, '('),
        Token(Tokens.NUM, '1'),
        Token(Tokens.PLUS, '+'),
        Token(Tokens.NUM, '2'),
        Token(Tokens.RPAREN, ')'),
        Token(Tokens.MUL, '*'),
        Token(Tokens.NUM, '3'),
    ]
    assert result == expected