diff options
Diffstat (limited to 'util')
| -rw-r--r-- | util/ast_extensions.py | 11 | ||||
| -rw-r--r-- | util/codegen.py | 14 |
2 files changed, 23 insertions, 2 deletions
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') |