| 58 | changed = False |
| 59 | |
| 60 | class FunctorTransformer(ast.NodeTransformer): |
| 61 | def __init__(self, target_functors): |
| 62 | super().__init__() |
| 63 | self.target_functors = target_functors |
| 64 | |
| 65 | # visit func def |
| 66 | def visit_FunctionDef(self, node: ast.FunctionDef): |
| 67 | if skip_main_fn and node.name == "main": |
| 68 | return node |
| 69 | # visit child nodes |
| 70 | return self.generic_visit(node) |
| 71 | |
| 72 | def visit_Call(self, node): |
| 73 | if isinstance(node.func, ast.Name) and node.func.id in self.target_functors: |
| 74 | nonlocal changed |
| 75 | changed = True |
| 76 | return ast.Expr(value=ast.Constant(value=None)) |
| 77 | return node |
| 78 | |
| 79 | tree = ast.parse(source) |
| 80 | code = astor.to_source(FunctorTransformer(functors).visit(tree)) |