fix export of 'f(x)' block

4 files changed, 27 insertions, 6 deletions

diff --git a/TurtleArt/taexportpython.py b/TurtleArt/taexportpython.py
index a97fbca..33c3979 100644
--- a/TurtleArt/taexportpython.py
+++ b/TurtleArt/taexportpython.py
@@ -39,9 +39,9 @@ from tautils import (debug_output, find_group, find_top_block, get_stack_name)
 _SETUP_CODE_START = """\
 #!/usr/bin/env python
 
-from math import sqrt
+from time import *
 from random import uniform
-from time import (sleep, time)
+from math import *
 
 from pyexported.window_setup import *
 
diff --git a/TurtleArt/taprimitive.py b/TurtleArt/taprimitive.py
index 7f30071..05a44ad 100644
--- a/TurtleArt/taprimitive.py
+++ b/TurtleArt/taprimitive.py
@@ -422,8 +422,8 @@ class Primitive(object):
             param_asts = []
             for id_ in ['x', 'y', 'z'][:len(new_arg_asts)-1]:
                 param_asts.append(ast.Name(id=id_, ctx=ast.Param))
-            func_ast = ast.Lambda(body=new_arg_asts[0], args=ast.arguments(
-                args=param_asts, vararg=None, kwarg=None, defaults=[]))
+            func_ast = ast_extensions.LambdaWithStrBody(
+                body_str=new_arg_asts[0].s, args=param_asts)
             return get_call_ast(func_ast, new_arg_asts[1:],
                 return_type=self.return_type)
 
diff --git a/util/ast_extensions.py b/util/ast_extensions.py
index ac29421..3335afb 100644
--- a/util/ast_extensions.py
+++ b/util/ast_extensions.py
@@ -39,3 +39,14 @@ class Comment(ast.stmt):
         self.col_offset = col_offset
 
 
+class LambdaWithStrBody(ast.Lambda):
+    """ Lambda AST whose body is a simple string (not ast.Str).
+    Extends the Python abstract grammar by the following:
+    expr = LambdaWithStrBody(string body_str, expr* args) | ... """
+
+    def __init__(self, body_str="", args=[], lineno=1, col_offset=0):
+        self.body_str = body_str
+        self.args = args
+        self.lineno = lineno
+        self.col_offset = col_offset
+
diff --git a/util/codegen.py b/util/codegen.py
index cd1d106..3785085 100644
--- a/util/codegen.py
+++ b/util/codegen.py
@@ -504,11 +504,21 @@ class SourceGenerator(NodeVisitor):
         self.visit(node.value)
 
     def visit_Lambda(self, node):
-        self.write('lambda ')
+        self.write('(lambda ')
         self.signature(node.args)
         self.write(': ')
         self.visit(node.body)
-    visit_TypedLambda = visit_Lambda
+        self.write(')')
+
+    def visit_LambdaWithStrBody(self, node):
+        self.write('(lambda ')
+        for idx, arg in enumerate(node.args):
+            if idx:
+                self.write(', ')
+            self.visit(arg)
+        self.write(': ')
+        self.write(node.body_str)
+        self.write(')')
 
     def visit_Ellipsis(self, node):
         self.write('Ellipsis')