-
-
Save rebcabin/ae7832088b94db5506d439f1fa5b451d to your computer and use it in GitHub Desktop.
just magic
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import sys | |
import ast | |
origin_import = __import__ | |
AST = {} | |
def custom_import(name, *args, **kwargs): | |
module = origin_import(name, *args, **kwargs) | |
if not hasattr(module, '__file__'): | |
return module | |
try: | |
mod_ast = ast.parse( | |
''.join( | |
open( | |
sys.modules[name].__file__.replace('pyc', 'py'), | |
'r' | |
).readlines() | |
) | |
) | |
patched_ast = patch(mod_ast) | |
exec(compile(patched_ast, name, "exec"), module.__dict__) | |
AST[name] = mod_ast | |
AST['%s_patched' % name] = patched_ast | |
return module | |
except: | |
return module | |
def gen_assign_checker_ast(targets, obj_name): | |
return ast.If( | |
test=ast.Call( | |
func=ast.Name(id='hasattr', ctx=ast.Load()), | |
args=[ | |
ast.Name(id=obj_name, ctx=ast.Load()), | |
ast.Str(s='__assign__'), | |
], | |
keywords=[], | |
starargs=None, | |
kwargs=None | |
), | |
body=[ | |
ast.Assign( | |
targets=[ast.Name(id=target, ctx=ast.Store())], | |
value=ast.Call( | |
func=ast.Attribute( | |
value=ast.Name(id=obj_name, ctx=ast.Load()), | |
attr='__assign__', | |
ctx=ast.Load() | |
), | |
args=[ast.Str(s=target)], | |
keywords=[], | |
starargs=None, | |
kwargs=None | |
) | |
) | |
for target in targets], | |
orelse=[] | |
) | |
class Transformer(ast.NodeTransformer): | |
def generic_visit(self, node): | |
ast.NodeTransformer.generic_visit(self, node) | |
return node | |
def visit_Assign(self, node): | |
if not isinstance(node.value, ast.Name): | |
return node | |
targets = [t.id for t in node.targets] | |
obj_name = node.value.id | |
new_node = gen_assign_checker_ast(targets, obj_name) | |
ast.copy_location(new_node, node) | |
ast.fix_missing_locations(new_node) | |
return new_node | |
def patch(node): | |
trans = Transformer() | |
new_node = trans.visit(node) | |
ast.fix_missing_locations(new_node) | |
return new_node | |
__builtins__.update(**dict( | |
__import__=custom_import, | |
__ast__=AST | |
)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
a = 1 | |
class T(): | |
def __assign__(self, v): | |
print('called with %s' % v) | |
b = T() | |
c = b |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment