tokin/parser.py

import dataclasses
import unittest

from lexer import Token, kinds, lex

class Expression: pass

@dataclasses.dataclass
class Atom(Expression):
	value: Token

@dataclasses.dataclass
class BinOp(Expression):
	operator: Token
	left: Expression
	right: Expression

@dataclasses.dataclass
class UnaryOp(Expression):
	operator: Token
	inner: Expression

class ParseError(Exception):
	def __init__(self, token, expected):
		Exception.__init__(self)
		self.token = token
		self.expected = expected

	def __repr__(self):
		return f'ParseError({repr(self.token)}, {repr(self.expected)})'

	def __str__(self):
		return self.__repr__()

def parse(string):
	tokens = list(lex(string))

	def peek():
		return tokens[0]

	def consume():
		return tokens.pop(0)

	def expect(kind, alternatives=None):
		if peek().kind != kind:
			if alternatives is None:
				alternatives = []
			raise ParseError(
				token=peek(),
				expected=[kind, *alternatives]
			)

	def binOpExpression(operators, inner):
		def perLevelLogic():
			tree = inner()
			while peek().kind in operators:
				operator = consume()
				tree = BinOp(
					operator=operator,
					left=tree,
					right=inner()
				)
			return tree
		return perLevelLogic

	def atom():
		if peek().kind == kinds.openparen:
			consume()
			tree = expression0()
			expect(kinds.closeparen)
			consume()
		else:
			expect(kinds.number, [kinds.openparen])
			tree = Atom(consume())
		return tree

	def expression2():
		if peek().kind == kinds.minus:
			operator = consume()
			return UnaryOp(
				operator=operator,
				inner=atom()
			)
		else:
			return atom()

	expression1 = binOpExpression((kinds.multiply, kinds.divide), expression2)
	expression0 = binOpExpression((kinds.plus, kinds.minus), expression1)

	tree = expression0()
	expect(kinds.eof)

	return tree

class Tests(unittest.TestCase):
	def test_addSub(self):
		self.assertEqual(parse('1+2-3-4+5'),
			BinOp(Token(7, '+', kinds.plus),
				BinOp(Token(5, '-', kinds.minus),
					BinOp(Token(3, '-', kinds.minus),
						BinOp(Token(1, '+', kinds.plus),
							Atom(Token(0, '1', kinds.number)),
							Atom(Token(2, '2', kinds.number))
						),
						Atom(Token(4, '3', kinds.number))
					),
					Atom(Token(6, '4', kinds.number))
				),
				Atom(Token(8, '5', kinds.number))
			)
		)

	def test_precedence(self):
		self.assertEqual(parse('1*2+3'),
			BinOp(Token(3, '+', kinds.plus),
				BinOp(Token(1, '*', kinds.multiply),
					Atom(Token(0, '1', kinds.number)),
					Atom(Token(2, '2', kinds.number))
				),
				Atom(Token(4, '3', kinds.number))
			)
		)
		self.assertEqual(parse('1-2/3'),
			BinOp(Token(1, '-', kinds.minus),
				Atom(Token(0, '1', kinds.number)),
				BinOp(Token(3, '/', kinds.divide),
					Atom(Token(2, '2', kinds.number)),
					Atom(Token(4, '3', kinds.number))
				)
			)
		)

	def test_parentheses(self):
		self.assertEqual(parse('1*(2+3)'),
			BinOp(Token(1, '*', kinds.multiply),
				Atom(Token(0, '1', kinds.number)),
				BinOp(Token(4, '+', kinds.plus),
					Atom(Token(3, '2', kinds.number)),
					Atom(Token(5, '3', kinds.number))
				)
			)
		)
		self.assertEqual(parse('(1-2)/3'),
			BinOp(Token(5, '/', kinds.divide),
				BinOp(Token(2, '-', kinds.minus),
					Atom(Token(1, '1', kinds.number)),
					Atom(Token(3, '2', kinds.number))
				),
				Atom(Token(6, '3', kinds.number))
			)
		)

	def test_negation(self):
		self.assertEqual(parse('-1*-1'),
			BinOp(Token(2, '*', kinds.multiply),
				UnaryOp(Token(0, '-', kinds.minus),
					Atom(Token(1, '1', kinds.number))
				),
				UnaryOp(Token(3, '-', kinds.minus),
					Atom(Token(4, '1', kinds.number))
				)
			)
		)
		self.assertEqual(parse('-(1+2)'),
			UnaryOp(Token(0, '-', kinds.minus),
				BinOp(Token(3, '+', kinds.plus),
					Atom(Token(2, '1', kinds.number)),
					Atom(Token(4, '2', kinds.number))
				)
			)
		)

	def test_parseErrors(self):
		with self.assertRaises(ParseError) as context:
			parse('2*)3+4)')
		self.assertEqual(
			context.exception.token,
			Token(2, ')', kinds.closeparen)
		)

		with self.assertRaises(ParseError) as context:
			parse('2*(3+4')
		self.assertEqual(
			context.exception.token,
			Token(6, '', kinds.eof)
		)

		with self.assertRaises(ParseError) as context:
			parse('(1+1)(2+2)')
		self.assertEqual(
			context.exception.token,
			Token(5, '(', kinds.openparen)
		)