add return type annotations (just decoration now)

3 files changed, 14 insertions, 8 deletions

diff --git a/neb/__init__.py b/neb/__init__.py
index f5574d1..4ae32a6 100644
--- a/neb/__init__.py
+++ b/neb/__init__.py
@@ -47,13 +47,14 @@ class Function:
         self.args = args
         self.many = many
         self.type_ = TypeEnum.ANY  # TODO no it's not
+        self.return_type = Type(":any")
 
     def describe(self, name=None):
         if name is None:
             name = self.name
         out = [f"({name}"]
         out.append(string_args(self.args, self.many))
-        return " ".join(out).strip() + ")"
+        return " ".join(out).strip() + f") {self.return_type}"
 
     def arity_check(self, symbol, params):
         min_arity = len([a for a in self.args if not a.optional])
@@ -115,7 +116,7 @@ class UserFunction(Function):
         super().__init__(name, newparams, body, args, many)
 
     def __str__(self):
-        out = f"(func {self.name} ("
+        out = f"(func {self.name} {self.return_type} ("
         args_list = [f"{a.name} {a.type_}" for a in self.args]
         if self.many:
             args_list.append(f"{self.many.name} {self.many.type_}")
diff --git a/neb/std/core.py b/neb/std/core.py
index 83b164c..d9f512e 100644
--- a/neb/std/core.py
+++ b/neb/std/core.py
@@ -45,10 +45,17 @@ def interpretRedef(symbol, args, env, ns):
 CORE.register("redef", Builtin("redef", interpretRedef, [def_name_arg, def_val_arg]))
 
 def interpretLambda(symbol, args, env, ns):
-    if len(args[0].args) != 0:
-        func = UserFunction("<lambda>", args[0].args, args[1:])
+    new_args = args
+    return_type = Type(":any")
+    if isinstance(args[0], Type):
+        return_type = args[0]
+        new_args = args[1:]
+
+    if len(new_args[0].args) != 0:
+        func = UserFunction("<lambda>", new_args[0].args, new_args[1:])
     else:
-        func = UserFunction("<lambda>", [], args[1:])
+        func = UserFunction("<lambda>", [], new_args[1:])
+    func.return_type = return_type
     return func
 
 lambda_args_arg = Arg("args", TypeEnum.ANY, lazy=True)
diff --git a/neb/structs.py b/neb/structs.py
index fcfa430..f6bd3af 100644
--- a/neb/structs.py
+++ b/neb/structs.py
@@ -126,9 +126,7 @@ class Arg:
         self.lazy = lazy
 
     def __str__(self):
-        opt = "?" if self.optional else ""
-        lazy = "~" if self.lazy else ""
-        return f"{lazy}{opt}{self.name} {self.type_}"
+        return f"{self.name} {self.type_}"
 
 def string_args(args, many):
     out = [f"{arg}" for arg in args]