def visit_If(self, node):
is_main = (isinstance(node.test, ast.Compare)
and isinstance(node.test.left, ast.Name)
and node.test.left.id == "__name__"
and isinstance(node.test.ops[0], ast.Eq)
and isinstance(node.test.comparators[0], ast.Constant)
and node.test.comparators[0].value == "__main__")
if is_main:
if hasattr(node, "scopes") and len(node.scopes) > 1:
rename(node.scopes[-2], "main", "main_func")
# ast.parse produces a Module object that needs to be destructured
if self.main_signature_arg_names == {"argc", "argv"}:
ret = cast(
ast.FunctionDef,
create_ast_node(
"def main(argc: int, argv: List[str]) -> int: True",
node),
)
elif self.main_signature_arg_names == {"argv"}:
ret = create_ast_node("def main(argv: List[str]): True", node)
else:
ret = create_ast_node("def main(): True")
ret = cast(ASTxFunctionDef, ret)
ret.lineno = node.lineno
ret.body = node.body
# So backends know to handle argc, argv etc
ret.python_main = True
return ret
return node
def visit_Call(self, node: ast.Call) -> ast.Call:
if (isinstance(node.func, ast.Attribute) and node.func.attr
== "values"): # and is_dict(node.func.value):
new_node: ast.Call = create_ast_node("a.keys().map(a[it])",
at_node=node)
new_node.func.value.func.value = node.func.value
new_node.args[0].value = node.func.value
return new_node
return node
def _do_other_rewrite(self, node) -> ast.AST:
ifexpr = cast(ast.Expr,
create_ast_node("'True' if true else 'False'",
node)).value
ifexpr = cast(ast.IfExp, ifexpr)
ifexpr.test = node.args[0]
ifexpr.lineno = node.lineno
ifexpr.col_offset = node.col_offset
ast.fix_missing_locations(ifexpr)
node.args[0] = ifexpr
return node
def visit_GeneratorExp(self, node: ast.GeneratorExp) -> ast.Call:
new_node = None
for comp in node.generators:
if isinstance(comp.target, ast.Name):
self.redirects[comp.target.id] = "it"
elif isinstance(comp.target, ast.Tuple):
for idx, elem in enumerate(comp.target.elts):
assert isinstance(elem, ast.Name)
self.redirects[elem.id] = f"it[{idx}]"
else:
raise AstNotImplementedError(
f"Unknown target type {type(node.target).__qualname__}",
node)
subnode = comp.iter
for cmp in comp.ifs:
chain = create_ast_node("placeholder.filter(placeholder)",
at_node=node)
chain.func.value = subnode
chain.args[0] = cmp
subnode = chain
chain = create_ast_node("placeholder.map(placeholder)",
at_node=node)
chain.func.value = subnode
chain.args[0] = node.elt
subnode = chain
if new_node is None:
new_node = subnode
else:
new_node.args[0] = subnode
self.visit(new_node)
self.redirects.clear()
return new_node
def _do_go_rewrite(self, node) -> ast.AST:
if_stmt = create_ast_node(
textwrap.dedent("""\
if True:
print('True')
else:
print('False')
"""),
node,
)
if_stmt = cast(ast.If, if_stmt)
if_stmt.test = node.args[0]
if_stmt.lineno = node.lineno
if_stmt.col_offset = node.col_offset
ast.fix_missing_locations(if_stmt)
return if_stmt
def visit_JoinedStr(self, node):
new_node = cast(ast.Expr, create_ast_node('"".join([])', node)).value
new_node = cast(ast.Call, new_node)
args = new_node.args
arg0 = cast(ast.List, args[0])
for v in node.values:
if isinstance(v, ast.Constant):
arg0.elts.append(v)
elif isinstance(v, ast.FormattedValue):
arg0.elts.append(
ast.Call(func=ast.Name(id="str", ctx="Load"),
args=[v.value],
keywords=[]))
new_node.lineno = node.lineno
new_node.col_offset = node.col_offset
ast.fix_missing_locations(new_node)
return new_node
def _annotate(node, typename: str):
# ast.parse produces a Module object that needs to be destructured
type_annotation = cast(ast.Expr, create_ast_node(typename, node)).value
node.annotation = type_annotation
def visit_Name(self, node: ast.Name) -> Union[ast.Name, ast.Subscript]:
if node.id in self.redirects:
return create_ast_node(self.redirects[node.id], at_node=node)
return node
