fix export of 'f(x)' block

1 files changed, 12 insertions, 2 deletions

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