introduce TypedLambda, a Lambda AST with a return type

1 files changed, 20 insertions, 0 deletions

diff --git a/TurtleArt/tatype.py b/TurtleArt/tatype.py
index 9cb5433..add26e3 100644
--- a/TurtleArt/tatype.py
+++ b/TurtleArt/tatype.py
@@ -342,6 +342,26 @@ class TypedCall(ast.Call):
             return self._return_type
 
 
+class TypedLambda(ast.Lambda):
+    """ Like a Lambda AST, but with a return type """
+
+    def __init__(self, body=None, args=None, return_type=None):
+
+        if args is None:
+            args = ast.arguments(args=[], vararg=None, kwarg=None, defaults=[])
+
+        ast.Lambda.__init__(self, body=body, args=args)
+
+        self._return_type = return_type
+
+    @property
+    def return_type(self):
+        if self._return_type is None:
+            return get_type(self.func)
+        else:
+            return self._return_type
+
+
 def get_call_ast(func_name, args=None, kwargs=None, return_type=None):
     """ Return an AST representing the call to a function with the name
     func_name, passing it the arguments args (given as a list) and the