✨

plthy_impl/ast_nodes.py

@@ -2,6 +2,13 @@ import random
 
 from rply.token import BaseBox
 
+
+BUILTIN_ARGS = {
+    "print": "*",
+    "input" : "*",
+    "random": 2
+}
+
 def rep_join(l):
     format_string = ',\n'.join(
         [repr(i) if not isinstance(i, str) else i for i in l]
@@ -11,7 +18,6 @@ def rep_join(l):
         format_string = f"\n  {format_string}\n"
 
     return format_string
-
 class Exp(BaseBox):
     def eval(self, vtable, ftable):
         return vtable, ftable, None
@@ -86,8 +92,7 @@ class ExpABinop(Exp):
         elif self.op == "<":
             return vtable, ftable, r1 < r2
         else:
-            print(f"E007: Unknown binop {self.op}")
-            exit()
+            raise Exception(f"Unknown binop {self.op}")
 
 class ExpList(Exp):
     def __init__(self, expressions: list[Exp]):
@@ -143,9 +148,15 @@ class Command(BaseBox):
         return f"command()"
 
 class Builtin(Command):
-    def __init__(self, builtin: str, args: list[Exp]):
-        self.builtin = builtin
+    def __init__(self, builtin, args: list[Exp]):
+        self.builtin = builtin.value
         self.args = args
+        self.position = builtin.source_pos.lineno
+        expected_args = BUILTIN_ARGS[self.builtin]
+
+        if not (expected_args == "*" or expected_args == len(self.args)):
+            print(f"E005: Wrong number of arguments for builtin '{self.builtin}' ({len(self.args)}) at line {self.position}")
+            exit()
 
     def __repr__(self) -> str:
         return f"builtin({self.builtin}, {self.args})"
@@ -174,10 +185,6 @@ class Builtin(Command):
             result = input()
             return vtable, ftable, result
         elif self.builtin == "random":
-            if not len(self.args) == 2:
-                print(f"E006: Wrong number of arguments for builtin '{self.builtin}' ({len(self.args)})")
-                exit()
-
             vtable, ftable, r1 = self.args[0].eval(vtable, ftable)
             vtable, ftable, r2 = self.args[1].eval(vtable, ftable)
 
@@ -185,7 +192,7 @@ class Builtin(Command):
 
             return vtable, ftable, r
         else:
-            raise Exception(f"E005: Unknown builtin {self.builtin}")
+            raise Exception(f"Unknown builtin {self.builtin}")
 
 class Do(BaseBox):
     def __init__(self, command: Command):
@@ -309,6 +316,22 @@ class StatementIf(Statement):
         else:
             return vtable, ftable, result
 
+class StatementBecause(Statement):
+    def __init__(self, statement: Statement, condition):
+        self.statement = statement
+        self.condition = condition
+
+    def __repr__(self) -> str:
+        return f"if({self.statement}, {self.condition})"
+
+    def eval(self, vtable, ftable):
+        vtable, ftable, result = self.condition.eval(vtable, ftable)
+        if result:
+            return self.statement.eval(vtable, ftable)
+        else:
+            print("C001: Because assertion incorrect")
+            exit()
+
 class StatementUntil(Statement):
     def __init__(self, statement: Statement, condition):
         self.statement = statement