introduce TypedLambda, a Lambda AST with a return type

2 files changed, 23 insertions, 4 deletions

diff --git a/TurtleArt/taprimitive.py b/TurtleArt/taprimitive.py
index ca0e73a..6b0c515 100644
--- a/TurtleArt/taprimitive.py
+++ b/TurtleArt/taprimitive.py
@@ -31,7 +31,7 @@ from taturtle import (Turtle, Turtles)
 from tatype import (ACTION_AST, BOX_AST, convert, get_call_ast, get_converter,
                     get_type, is_bound_instancemethod, is_instancemethod,
                     is_staticmethod, TATypeError, Type, TypeDisjunction,
-                    TYPE_COLOR, TYPE_FLOAT, TYPE_OBJECT)
+                    TypedLambda, TYPE_COLOR, TYPE_FLOAT, TYPE_OBJECT)
 from tautils import debug_output
 from tawindow import (global_objects, TurtleArtWindow)
 from util import ast_extensions
@@ -850,9 +850,8 @@ class ArgSlot(object):
                             # don't call and pass on the callable
                             if convert_to_ast:
                                 lambda_body = func_prim.get_ast()
-                                called_argument = ast.Lambda(
-                                    body=lambda_body, args=[],
-                                    vararg=None, kwarg=None, defaults=[])
+                                called_argument = TypedLambda(body=lambda_body,
+                                    return_type=lambda_body.return_type)
                             else:
                                 called_argument = func_prim
 
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