refactor: take type hints from user in function args

5 files changed, 66 insertions, 40 deletions

diff --git a/interpreter.py b/interpreter.py
index 7f28c63..e66deab 100644
--- a/interpreter.py
+++ b/interpreter.py
@@ -2,6 +2,7 @@ from structs import *
 from exceptions import *
 from lexer import lex
 from parser import parse
+from typeclass import TypeEnum, is_subtype_of
 from pathlib import Path
 from glob import glob
 from collections import namedtuple
@@ -99,8 +100,26 @@ class UserFunction(Function):
     
     def __init__(self, name, params, body):
         # TODO this doesn't do type checking, or optional, or lazy
-        args = [Arg(p.name, TypeEnum.ANY, False, False) for p in params]
-        super().__init__(name, params, body, args)
+        newparams, args = self.process_params(name, params)
+        super().__init__(name, newparams, body, args)
+
+    def process_params(self, name, params):
+        newparams = []
+        args = []
+        prev_type = False
+        first = True
+        for param in params:
+            if isinstance(param, Symbol):
+                newparams.append(param)
+                args.append(Arg(param.name, TypeEnum.ANY, False, False))
+                prev_type = False
+            elif isinstance(param, Type) and not prev_type and not first:
+                args[-1].type_ = TypeEnum.__getattr__(param.name[1:].upper())
+                prev_type = True
+            else:
+                raise NebPanic("invalid :func signature", param)
+            first = False
+        return newparams, args
 
     def call(self, expr, env, ns):
         self.arity_check(expr.args[0], expr.args[1:])
@@ -108,7 +127,6 @@ class UserFunction(Function):
         this_env = Environment(env)
         for idx, param in enumerate(self.params):
             this_env.register(param.name, evaluated_args[idx])
-        #return interpret(self.body, this_env, ns)
         return interpret(self.body, env=this_env, ns=ns)
 
 class Environment:
@@ -159,7 +177,7 @@ def interpret(exprs, *, env=GLOBALS, ns=None):
     return ret
 
 def evaluate(expr, env, ns=None):
-    if isinstance(expr, Literal) or isinstance(expr, Function):
+    if isinstance(expr, Literal) or isinstance(expr, Function) or isinstance(expr, Type):
         #return expr.value
         return expr
     elif isinstance(expr, Symbol):
diff --git a/lexer.py b/lexer.py
index 0a634c5..ea4fefb 100644
--- a/lexer.py
+++ b/lexer.py
@@ -9,9 +9,12 @@ class LexError(BaseException):
 types = {
     ":int": TokenType.INT_TYPE,
     ":float": TokenType.FLOAT_TYPE,
+    ":number": TokenType.NUMBER_TYPE,
     ":string": TokenType.STRING_TYPE,
     ":list": TokenType.LIST_TYPE,
-    ":any": TokenType.ANY_TYPE }
+    ":any": TokenType.ANY_TYPE,
+    ":literal": TokenType.LITERAL_TYPE,
+    ":bool": TokenType.BOOL_TYPE }
 
 keywords = {
     "print": TokenType.PRINT,
diff --git a/parser.py b/parser.py
index 60bd029..46e43e7 100644
--- a/parser.py
+++ b/parser.py
@@ -13,7 +13,7 @@ def parseExpression(token, prev, tokens):
         elif token.type_ in (TokenType.STRING, TokenType.TRUE, TokenType.FALSE, TokenType.INT, TokenType.FLOAT):
             expr, inc = parseLiteral(token, prev, tokens[idx+1:])
             args.append(expr)
-        elif token.type_ in (TokenType.INT_TYPE, TokenType.FLOAT_TYPE, TokenType.STRING_TYPE, TokenType.ANY_TYPE):
+        elif token.type_ in (TokenType.INT_TYPE, TokenType.FLOAT_TYPE, TokenType.STRING_TYPE, TokenType.ANY_TYPE, TokenType.LIST_TYPE, TokenType.NUMBER_TYPE, TokenType.BOOL_TYPE, TokenType.LITERAL_TYPE):
             expr, inc = parseType(token, prev, tokens[idx+1:])
             args.append(expr)
         else:
@@ -55,7 +55,7 @@ def parse(tokens):
         elif token.type_ in (TokenType.FALSE, TokenType.TRUE, TokenType.STRING, TokenType.INT, TokenType.FLOAT):
             lit, counter = parseLiteral(token, prev, tokens[idx+1:])
             exprs.append(lit)
-        elif token.type_ in (TokenType.INT_TYPE, TokenType.FLOAT_TYPE, TokenType.STRING_TYPE, TokenType.ANY_TYPE):
+        elif token.type_ in (TokenType.INT_TYPE, TokenType.FLOAT_TYPE, TokenType.STRING_TYPE, TokenType.ANY_TYPE, TokenType.LIST_TYPE, TokenType.NUMBER_TYPE, TokenType.BOOL_TYPE, TokenType.LITERAL_TYPE):
             typ, counter = parseType(token, prev, tokens[idx+1:])
             exprs.append(typ)
         else:
diff --git a/structs.py b/structs.py
index 72fabf7..adfd0e4 100644
--- a/structs.py
+++ b/structs.py
@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from enum import Enum, auto
 from typing import Any
+from typeclass import TypeEnum
 
 # tokens and types
 # NOTE: this can probably be simplified
@@ -54,42 +55,12 @@ class TokenType(Enum):
     # types
     INT_TYPE = auto()
     FLOAT_TYPE = auto()
+    NUMBER_TYPE = auto()
     STRING_TYPE = auto()
     ANY_TYPE = auto()
     LIST_TYPE = auto()
-
-
-class TypeEnum(Enum):
-    ANY = auto()
-    STRING = auto()
-    INT = auto()
-    FLOAT = auto()
-    NUMBER = auto()
-    LIST = auto()
-    LITERAL = auto()
-    BOOL = auto()
-
-    def __str__(self):
-        return f":{self.name.lower()}"
-
-TYPE_HIERARCHY = { TypeEnum.ANY: None,
-                    TypeEnum.LITERAL: TypeEnum.ANY,
-                    TypeEnum.LIST: TypeEnum.ANY,
-                    TypeEnum.STRING: TypeEnum.LITERAL,
-                    TypeEnum.BOOL: TypeEnum.LITERAL,
-                    TypeEnum.NUMBER: TypeEnum.LITERAL,
-                    TypeEnum.INT: TypeEnum.NUMBER,
-                    TypeEnum.FLOAT: TypeEnum.NUMBER }
-
-def is_subtype_of(candidate, expected):
-    if candidate == expected:
-        return True
-    parent = TYPE_HIERARCHY[candidate]
-    while parent is not None:
-        if parent == expected:
-            return True
-        parent = TYPE_HIERARCHY[parent]
-    return False
+    LITERAL_TYPE = auto()
+    BOOL_TYPE = auto()
 
 @dataclass
 class Token:
diff --git a/typeclass.py b/typeclass.py
new file mode 100644
index 0000000..eae412c
--- /dev/null
+++ b/typeclass.py
@@ -0,0 +1,34 @@
+from enum import Enum, auto
+
+class TypeEnum(Enum):
+    ANY = auto()
+    STRING = auto()
+    INT = auto()
+    FLOAT = auto()
+    NUMBER = auto()
+    LIST = auto()
+    LITERAL = auto()
+    BOOL = auto()
+
+    def __str__(self):
+        return f":{self.name.lower()}"
+
+HIERARCHY = { TypeEnum.ANY: None,
+                    TypeEnum.LITERAL: TypeEnum.ANY,
+                    TypeEnum.LIST: TypeEnum.ANY,
+                    TypeEnum.STRING: TypeEnum.LITERAL,
+                    TypeEnum.BOOL: TypeEnum.LITERAL,
+                    TypeEnum.NUMBER: TypeEnum.LITERAL,
+                    TypeEnum.INT: TypeEnum.NUMBER,
+                    TypeEnum.FLOAT: TypeEnum.NUMBER }
+
+def is_subtype_of(candidate, expected):
+    if candidate == expected:
+        return True
+    parent = HIERARCHY[candidate]
+    while parent is not None:
+        if parent == expected:
+            return True
+        parent = HIERARCHY[parent]
+    return False
+