make documentation look nicer

1 files changed, 4 insertions, 2 deletions

diff --git a/neb/__init__.py b/neb/__init__.py
index 4ae32a6..9b3c1e9 100644
--- a/neb/__init__.py
+++ b/neb/__init__.py
@@ -54,7 +54,7 @@ class Function:
             name = self.name
         out = [f"({name}"]
         out.append(string_args(self.args, self.many))
-        return " ".join(out).strip() + f") {self.return_type}"
+        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])
@@ -97,8 +97,10 @@ class Function:
 
 class Builtin(Function):
     
-    def __init__(self, name, callable_, args=None, many=None):
+    def __init__(self, name, callable_, args=None, many=None, return_type=None):
         super().__init__(name, None, callable_, args, many)
+        if return_type is not None:
+            self.return_type = return_type
 
     def __str__(self):
         return f"builtin function {self.name}"