introduce TypedCall, a Call AST that knows its return type

  - Add a visitor for TypedCall to util.codegen.

3 files changed, 65 insertions, 28 deletions

diff --git a/TurtleArt/taprimitive.py b/TurtleArt/taprimitive.py
index a5597d0..825009a 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_FLOAT, TYPE_OBJECT)
+                    TYPE_COLOR, TYPE_FLOAT, TYPE_OBJECT)
 from tautils import debug_output
 from tawindow import (global_objects, TurtleArtWindow)
 from util import ast_extensions
@@ -377,7 +377,8 @@ class Primitive(object):
                     return get_type(x)[0] == TYPE_FLOAT
                 if (    not _is_float(new_arg_asts[0]) and
                         not _is_float(new_arg_asts[1])):
-                    new_arg_asts[0] = get_call_ast('float', [new_arg_asts[0]])
+                    new_arg_asts[0] = get_call_ast('float', [new_arg_asts[0]],
+                        return_type=TYPE_FLOAT)
             if len(new_arg_asts) == 1:
                 if isinstance(op, tuple):
                     op = op[0]
@@ -402,7 +403,8 @@ class Primitive(object):
 
         # square root
         elif self == Primitive.square_root:
-            return get_call_ast('sqrt', new_arg_asts, new_kwarg_asts)
+            return get_call_ast('sqrt', new_arg_asts, new_kwarg_asts,
+                return_type=self.return_type)
 
         # identity
         elif self == Primitive.identity:
@@ -432,7 +434,8 @@ class Primitive(object):
 
         else:
             func_name = self.get_name_for_export()
-            return get_call_ast(func_name, new_arg_asts, new_kwarg_asts)
+            return get_call_ast(func_name, new_arg_asts, new_kwarg_asts,
+                return_type=self.return_type)
 
     def __eq__(self, other):
         """ Two Primitives are equal iff their all their properties are equal.
@@ -991,7 +994,8 @@ def value_to_ast(value, *args_for_prim, **kwargs_for_prim):
             # call to the Color constructor with this object's values,
             # e.g., Color('red', 0, 50, 100)
             return get_call_ast('Color', [value.name, value.color,
-                                          value.shade, value.gray])
+                                          value.shade, value.gray],
+                                return_type=TYPE_COLOR)
     else:
         raise ValueError("unknown type of raw value: " + repr(type(value)))
 
diff --git a/TurtleArt/tatype.py b/TurtleArt/tatype.py
index b2b14f4..9cb5433 100644
--- a/TurtleArt/tatype.py
+++ b/TurtleArt/tatype.py
@@ -183,24 +183,6 @@ TYPE_CONVERTERS = {
 }
 
 
-def get_call_ast(func_name, args=None, kwargs=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
-    keyword arguments kwargs (given as a dictionary). """
-    if args is None:
-        args = []
-    keywords = []
-    if kwargs is not None:
-        for (key, value) in kwargs.iteritems():
-            keywords.append(ast.keyword(arg=key, value=value))
-    return ast.Call(func=ast.Name(id=func_name,
-                                  ctx=ast.Load),
-                    args=args,
-                    keywords=keywords,
-                    starargs=None,
-                    kwargs=None)
-
-
 class TATypeError(BaseException):
     """ TypeError with the types from the hierarchy, not with Python types """
 
@@ -319,11 +301,7 @@ def convert(x, new_type, old_type=None, converter=None):
                 func = ast.Attribute(value=y,
                                      attr=converter.im_func.__name__,
                                      ctx=ast.Load)
-                return ast.Call(func=func,
-                                args=[],
-                                keywords={},
-                                starargs=None,
-                                kwargs=None)
+                return get_call_ast(func)
             else:
                 func_name = converter.__name__
                 return get_call_ast(func_name, [y])
@@ -340,3 +318,57 @@ def convert(x, new_type, old_type=None, converter=None):
         return _apply_converter(converter, x)
 
 
+class TypedCall(ast.Call):
+    """ Like a Call AST, but with a return type """
+
+    def __init__(self, func, args=None, keywords=None, starargs=None,
+            kwargs=None, return_type=None):
+
+        if args is None:
+            args = []
+        if keywords is None:
+            keywords = []
+
+        ast.Call.__init__(self, func=func, args=args, keywords=keywords,
+            starargs=starargs, kwargs=kwargs)
+
+        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
+    keyword arguments kwargs (given as a dictionary).
+    func_name -- either the name of a callable as a string, or an AST
+        representing a callable expression
+    return_type -- if this is not None, return a TypedCall object with this
+        return type instead """
+    if args is None:
+        args = []
+    # convert keyword argument dict to a list of (key, value) pairs
+    keywords = []
+    if kwargs is not None:
+        for (key, value) in kwargs.iteritems():
+            keywords.append(ast.keyword(arg=key, value=value))
+    # get or generate the AST representing the callable
+    if isinstance(func_name, ast.AST):
+        func_ast = func_name
+    else:
+        func_ast = ast.Name(id=func_name, ctx=ast.Load)
+    # if no return type is given, return a simple Call AST
+    if return_type is None:
+        return ast.Call(func=func_ast, args=args, keywords=keywords,
+            starargs=None, kwargs=None)
+    # if a return type is given, return a TypedCall AST
+    else:
+        return TypedCall(func=func_ast, args=args, keywords=keywords,
+            return_type=return_type)
+
+
diff --git a/util/codegen.py b/util/codegen.py
index 1bcc42f..f892308 100644
--- a/util/codegen.py
+++ b/util/codegen.py
@@ -393,6 +393,7 @@ class SourceGenerator(NodeVisitor):
             self.write('**')
             self.visit(node.kwargs)
         self.write(')')
+    visit_TypedCall = visit_Call
 
     def visit_Name(self, node):
         self.write(node.id)