introduce TypedCall, a Call AST that knows its return type

  - Add a visitor for TypedCall to util.codegen.

1 files changed, 55 insertions, 23 deletions

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)
+
+