<s> import ast <EOL> def is_write ( node ) : <EOL> return isinstance ( node , ( ast . Import , ast . ImportFrom , <EOL> ast . FunctionDef , ast . ClassDef , ast . arguments ) ) or isinstance ( node . ctx , ( ast . Store , ast . Del , ast . Param ) ) <EOL> def is_use ( node ) : <EOL> return isinstance ( node , ast . Name ) and isinstance ( node . ctx , ( ast . Load , ast . Del ) ) <EOL> def is_constant ( node ) : <EOL> return isinstance ( node , ast . Name ) and node . id . isupper ( ) <EOL> def is_param ( node ) : <EOL> return isinstance ( node , ast . Name ) and isinstance ( node . ctx , ast . Param ) or isinstance ( node , ast . arguments ) <EOL> def is_import ( node ) : <EOL> return isinstance ( node , ( ast . Import , ast . ImportFrom ) ) <EOL> def is_name ( node , name ) : <EOL> return isinstance ( node , ast . Name ) and node . id == name <EOL> def ast_eval ( node ) : <EOL> if isinstance ( node , ( ast . List , ast . Tuple ) ) : <EOL> return map ( ast_eval , node . elts ) <EOL> elif isinstance ( node , ast . Str ) : <EOL> return node . s <EOL> elif isinstance ( node , ast . Num ) : <EOL> return node . n <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % node . __class__ . __name__ ) <EOL> def name_class ( node ) : <EOL> if isinstance ( node , ( ast . Import , ast . ImportFrom ) ) : <EOL> return '<STR_LIT>' <EOL> elif isinstance ( node , ast . FunctionDef ) : <EOL> return '<STR_LIT>' <EOL> elif isinstance ( node , ast . ClassDef ) : <EOL> return '<STR_LIT:class>' <EOL> elif is_param ( node ) : <EOL> return '<STR_LIT>' <EOL> else : <EOL> return '<STR_LIT>' <EOL> def node_str ( node ) : <EOL> return '<STR_LIT>' % ( name_class ( node ) , node . lineno , node . col_offset ) <EOL> def nodes_str ( nodes ) : <EOL> return '<STR_LIT>' % '<STR_LIT:U+002CU+0020>' . join ( map ( node_str , nodes ) ) <EOL> import sys <EOL> import inspect <EOL> import textwrap <EOL> def get_body_ast ( func ) : <EOL> return get_ast ( func ) . body [ <NUM_LIT:0> ] . body <EOL> def get_ast ( func ) : <EOL> source = inspect . getsource ( func ) <EOL> source = textwrap . dedent ( source ) <EOL> source = '<STR_LIT:\n>' * ( func . __code__ . co_firstlineno - <NUM_LIT:2> ) + source <EOL> return ast . parse ( source , func_file ( func ) , '<STR_LIT>' ) <EOL> def func_file ( func ) : <EOL> return getattr ( sys . modules [ func . __module__ ] , '<STR_LIT>' , '<STR_LIT>' ) <EOL> from astor . codegen import SourceGenerator <EOL> from termcolor import colored <EOL> def to_source ( node , indent_with = '<STR_LIT:U+0020>' * <NUM_LIT:4> , add_line_information = False ) : <EOL> """<STR_LIT>""" <EOL> generator = AnnotatedSourceGenerator ( indent_with , add_line_information ) <EOL> generator . visit ( node ) <EOL> return '<STR_LIT>' . join ( str ( s ) for s in generator . result ) <EOL> class AnnotatedSourceGenerator ( SourceGenerator ) : <EOL> def visit ( self , node ) : <EOL> SourceGenerator . visit ( self , node ) <EOL> if not isinstance ( node , ( ast . Num , ast . Str ) ) and hasattr ( node , '<STR_LIT>' ) : <EOL> self . write ( colored ( '<STR_LIT>' % node . val , '<STR_LIT>' ) ) </s>
