<s> import ast <EOL> import inspect <EOL> from funcy import zipdict , is_list , keep <EOL> import astor <EOL> from . asttools import get_body_ast <EOL> class Pattern ( object ) : <EOL> pass <EOL> class UselessIf ( Pattern ) : <EOL> def template ( cond = ast . expr ) : <EOL> if cond : <EOL> return True <EOL> else : <EOL> return False <EOL> def suggestion ( ) : <EOL> return bool ( cond ) <EOL> class MapLambda ( Pattern ) : <EOL> def template ( body = ast . expr , seq = ast . expr ) : <EOL> map ( lambda var : body , seq ) <EOL> def suggestion ( ) : <EOL> [ body for var in seq ] <EOL> def match ( template , tree ) : <EOL> stack = [ ] <EOL> potential = [ ] <EOL> matches = [ ] <EOL> def _match ( node ) : <EOL> next_potential = [ ] <EOL> for p in potential : <EOL> if stack [ : len ( p [ '<STR_LIT>' ] ) ] != p [ '<STR_LIT>' ] : <EOL> matches . append ( p ) <EOL> else : <EOL> path = stack [ len ( p [ '<STR_LIT>' ] ) : ] <EOL> sub_template = get_sub_template ( template , path ) <EOL> if node_matches ( node , sub_template , p [ '<STR_LIT>' ] ) : <EOL> next_potential . append ( p ) <EOL> potential [ : ] = next_potential <EOL> context = { '<STR_LIT>' : { } , '<STR_LIT>' : { } , '<STR_LIT>' : { } } <EOL> if node_matches ( node , template , context ) : <EOL> potential . append ( { <EOL> '<STR_LIT>' : stack [ : ] , <EOL> '<STR_LIT>' : node [ <NUM_LIT:0> ] if is_list ( node ) else node , <EOL> '<STR_LIT>' : context , <EOL> } ) <EOL> if isinstance ( node , ast . AST ) : <EOL> for name , value in ast . iter_fields ( node ) : <EOL> stack . append ( name ) <EOL> _match ( value ) <EOL> stack . pop ( ) <EOL> elif isinstance ( node , list ) and node : <EOL> stack . append ( <NUM_LIT:0> ) <EOL> _match ( node [ <NUM_LIT:0> ] ) <EOL> stack . pop ( ) <EOL> stack . append ( <NUM_LIT:1> ) <EOL> _match ( node [ <NUM_LIT:1> : ] ) <EOL> stack . pop ( ) <EOL> _match ( tree ) <EOL> results = [ ] <EOL> for m in matches + potential : <EOL> m [ '<STR_LIT>' ] . captures = m [ '<STR_LIT>' ] [ '<STR_LIT>' ] <EOL> results . append ( m [ '<STR_LIT>' ] ) <EOL> return results <EOL> def node_matches ( node , template_node , context ) : <EOL> if isinstance ( template_node , ast . AST ) : <EOL> return type ( node ) is type ( template_node ) <EOL> elif isinstance ( template_node , list ) : <EOL> return isinstance ( node , list ) and len ( node ) >= len ( template_node ) and ( template_node == [ ] or node_matches ( node [ <NUM_LIT:0> ] , template_node [ <NUM_LIT:0> ] , context ) ) <EOL> elif template_node is None : <EOL> return node is None <EOL> elif isinstance ( template_node , ( str , int , float ) ) : <EOL> return node == template_node <EOL> else : <EOL> return template_node ( node , context ) <EOL> def get_sub_template ( template , path ) : <EOL> sub = template <EOL> for el in path : <EOL> if el == <NUM_LIT:0> : <EOL> try : <EOL> sub = sub [ <NUM_LIT:0> ] <EOL> except IndexError : <EOL> return lambda node , _ : False <EOL> elif el == <NUM_LIT:1> : <EOL> sub = sub [ <NUM_LIT:1> : ] <EOL> elif isinstance ( sub , ast . AST ) and el in sub . _fields : <EOL> sub = getattr ( sub , el ) <EOL> elif callable ( sub ) : <EOL> return lambda node , _ : True <EOL> else : <EOL> raise Exception ( '<STR_LIT>' , path , '<STR_LIT>' , astor . dump ( sub ) ) <EOL> return sub <EOL> def compile_template ( func ) : <EOL> spec = inspect . getargspec ( func ) <EOL> assert len ( spec . args ) == len ( spec . defaults or [ ] ) , "<STR_LIT>" <EOL> compiler = TemplateCompiler ( zipdict ( spec . args , spec . defaults or [ ] ) ) <EOL> template = map ( compiler . visit , get_body_ast ( func ) ) <EOL> if len ( template ) == <NUM_LIT:1> and isinstance ( template [ <NUM_LIT:0> ] , ast . Expr ) : <EOL> return template [ <NUM_LIT:0> ] . value <EOL> return template <EOL> class TemplateCompiler ( ast . NodeTransformer ) : <EOL> def __init__ ( self , args ) : <EOL> self . args = args <EOL> def generic_visit ( self , node ) : <EOL> """<STR_LIT>""" <EOL> for field , old_value in ast . iter_fields ( node ) : <EOL> old_value = getattr ( node , field , None ) <EOL> if isinstance ( old_value , list ) : <EOL> old_value [ : ] = keep ( self . visit , old_value ) <EOL> elif isinstance ( old_value , ast . AST ) : <EOL> new_node = self . visit ( old_value ) <EOL> if new_node is None : <EOL> delattr ( node , field ) <EOL> else : <EOL> setattr ( node , field , new_node ) <EOL> return node <EOL> def visit_Attribute ( self , node ) : <EOL> if isinstance ( node . value , ast . Name ) and node . value . id == '<STR_LIT>' : <EOL> cls = getattr ( ast , node . attr ) <EOL> return lambda n , _ : isinstance ( n , cls ) <EOL> else : <EOL> return node <EOL> def visit_Name ( self , node ) : <EOL> if node . id in { '<STR_LIT:True>' , '<STR_LIT:False>' , '<STR_LIT:None>' } : <EOL> return node <EOL> elif node . id in self . args : <EOL> def match_capture ( n , context ) : <EOL> if not isinstance ( n , self . args [ node . id ] ) : <EOL> return False <EOL> context [ '<STR_LIT>' ] [ node . id ] = n <EOL> if self . args [ node . id ] == ast . Name : <EOL> if ( node . id in context [ '<STR_LIT>' ] ) != ( n . id in context [ '<STR_LIT>' ] ) : <EOL> return False <EOL> if node . id in context [ '<STR_LIT>' ] : <EOL> return n . id == context [ '<STR_LIT>' ] [ node . id ] <EOL> else : <EOL> context [ '<STR_LIT>' ] [ node . id ] = n . id <EOL> context [ '<STR_LIT>' ] [ n . id ] = node . id <EOL> return True <EOL> return True <EOL> return match_capture <EOL> else : <EOL> return node </s>
