lisp_interpreter.py

import unittest

"""
Write code that takes some Lisp code and returns an abstract syntax tree.

The AST should represent the structure of the code and the meaning of each token.

For example, if your code is given "(first (list 1 (+ 2 3) 9))", it could return a nested array like
    ["first", ["list", 1, ["+", 2, 3], 9]].

Other examples:
    (+ 1 2) -> ["+", 1, 2]
    (* (+ 2 3) 9) -> ["*", parseAST("(+ 2 3)"), 9]
    (* 1 (+ 2 3)) -> ["*" 1, parseAST("(+ 2 3)")]
    (* (- 7 2) (+ 2 3)) -> ["*", parseAST("(- 7 2)"), parseAST("(+ 2 3)")]
    (* (+ (- 7 2) 4) 9) -> ["*", parseAST("(+ (-7 2) 4)"), 9]
"""
class LispInterpreter:
    def __init__(self):
        self.valid_operators = {"-", "+", "/", "*", "list", "first"}

    # parse a string representing some lisp code into an AST represented by nested arrays
    def parseAST(self, input: str) -> list[str | int]:
        results = []

        input = input.strip()

        if len(input) == 0:
            return results

        try:
            # validations that wouldn't cause an error
            if input[0] != "(":
                raise ValueError(f"expected '(' to start expression, got '{input}'")
            if input[-1] != ")":
                raise ValueError(f"expected ')' to start expression, got '{input}'")

            # strip open/close parens
            input = input[1:-1]
            # parse remaining elements
            # any parens get parsed recursively
            for token in self.nextToken(input):
                # convert numbers to int
                if token.isnumeric():
                    token = int(token)
                # if not a valid operator, assume next token is a nested expression
                elif token not in self.valid_operators:
                    try:
                        token = self.parseAST(token)
                    except:
                        raise ValueError(f"invalid operator or expression '{token}'")
                results.append(token)
        except Exception as e:
            raise ValueError(f"invalid expression received: '{input}' -> {e}")

        return results

    # generator that yields the next element or expression until input has been exhausted
    # NOTE: does not parse the token if it's an expression
    def nextToken(self, input: str) -> str:
        while input:
            # if first char is paren, find matching close paren
            if input[0] == "(":
                open_count = 1
                close_index = 0
                for i in range(1, len(input)):
                    if input[i] == "(":
                        open_count += 1
                    elif input[i] == ")":
                        open_count -= 1
                    if open_count == 0:
                        close_index = i
                        break
                next_token = input[: close_index + 1]
            # otherwise, return element up to next space
            else:
                next_token = input.split(" ")[0]

            token_end_index = len(next_token)
            input = input[token_end_index + 1 :]

            yield next_token


class TestParseAST(unittest.TestCase):
    def setUp(self):
        self.li = LispInterpreter()
        self.subject = self.li.parseAST

    def test_missing_lparen(self):
        input = "+ 2 3)"
        with self.assertRaises(ValueError) as e:
            self.subject(input)
        self.assertTrue(
            f"expected '(' to start expression, got '{input}'" in str(e.exception)
        )

    def test_missing_rparen(self):
        input = "(+ 2 3"
        with self.assertRaises(ValueError) as e:
            self.subject(input)
        self.assertTrue(
            f"expected ')' to start expression, got '{input}'" in str(e.exception)
        )

    # TODO: this is an assumption. need to clarify if this should actually be None
    def test_empty_input(self):
        input = ""
        self.assertEqual(self.subject(input), [])

    def test_invalid_input(self):
        input = "()"
        self.subject(input)

    def test_simple_expr(self):
        input = "(+ 2 3)"
        self.assertEqual(self.subject(input), ["+", 2, 3])

    def test_nested_expr(self):
        input = "(+ (- 3 4) 3)"
        self.assertEqual(self.subject(input), ["+", ["-", 3, 4], 3])

    def test_nested_expr_second_arg(self):
        input = "(+ 3 (- 3 4))"
        self.assertEqual(self.subject(input), ["+", 3, ["-", 3, 4]])

    def test_nested_expr_both_args(self):
        input = "(* (- 7 2) (+ 2 3))"
        self.assertEqual(self.subject(input), ["*", ["-", 7, 2], ["+", 2, 3]])

    def test_doubly_nested_expr(self):
        input = "(+ (- 3 (* 1 4)) 3)"
        self.assertEqual(self.subject(input), ["+", ["-", 3, ["*", 1, 4]], 3])

    def test_string_operator(self):
        input = "(list 1 2 3)"
        self.assertEqual(self.subject(input), ["list", 1, 2, 3])

    def test_invalid_operator(self):
        input = "(foo 2 3)"
        with self.assertRaises(ValueError) as e:
            self.subject(input)
        self.assertTrue("invalid operator or expression 'foo'" in str(e.exception))


class TestNextToken(unittest.TestCase):
    def setUp(self):
        self.li = LispInterpreter()
        self.subject = self.li.nextToken

    def test_happy_path(self):
        input = "+ 2 3"
        iter = self.subject(input)
        self.assertEqual(list(iter), ["+", "2", "3"])

    def test_empty_input(self):
        input = ""
        iter = self.subject(input)
        with self.assertRaises(StopIteration):
            next(iter)

    def test_one_arg(self):
        input = "+"
        iter = self.subject(input)
        self.assertEqual(list(iter), ["+"])

    def test_two_args(self):
        input = "+ 2"
        iter = self.subject(input)
        self.assertEqual(list(iter), ["+", "2"])

    def test_four_args(self):
        input = "+ 2 3 4"
        iter = self.subject(input)
        self.assertEqual(list(iter), ["+", "2", "3", "4"])

    def test_nested_second_arg(self):
        input = "+ (- 2 3) 4"
        iter = self.subject(input)
        self.assertEqual(next(iter), "+")
        self.assertEqual(next(iter), "(- 2 3)")

    def test_nested_third_arg(self):
        input = "+ 4 (- 2 3)"
        iter = self.subject(input)
        self.assertEqual(next(iter), "+")
        self.assertEqual(next(iter), "4")
        self.assertEqual(next(iter), "(- 2 3)")


if __name__ == "__main__":
    li = LispInterpreter()
    unittest.main()