def generic_visit(self, node): if (isinstance(node, ast.stmt) and self.references_arg(node)) or isinstance(node, ast.Return): return self.get_waits() + [node] return NodeTransformer.generic_visit(self, node)