Skip to content

Instantly share code, notes, and snippets.

@jargnar
Last active March 20, 2021 16:12
Show Gist options
  • Save jargnar/08564afd7298a34d3c84 to your computer and use it in GitHub Desktop.
Save jargnar/08564afd7298a34d3c84 to your computer and use it in GitHub Desktop.
Parse a Python AST and delete stuff from it
import ast
class RemoveMethod(ast.NodeTransformer):
'''
Removes all occurences of a method from the AST
Example:
---
tree = RemoveMethod('bar').visit(tree)
x = foo.bar() => x = foo
y = meow('hi').bar() => y = meow('hi')
'''
def __init__(self, func):
self.func = func
def visit_Call(self, node):
transform = node
try:
if node.func.attr == self.func:
if not len(node.args):
transform = node.func.value
except AttributeError:
pass
return ast.copy_location(transform, node)
tree = RemoveMethod('bar').visit(tree)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment