<s> import ast <EOL> import parser <EOL> import symbol <EOL> import sys <EOL> import token <EOL> class STTree ( object ) : <EOL> def __init__ ( self , s ) : <EOL> st = parser . suite ( s ) <EOL> st_tup = st . totuple ( ) <EOL> self . root = STNode ( st_tup ) <EOL> def find ( self , val , max_depth = None ) : <EOL> return self . root . find ( val , max_depth = max_depth ) <EOL> def find_all ( self , val , max_depth = None ) : <EOL> return self . root . find_all ( val , max_depth = max_depth ) <EOL> def find_bfs ( self , val , max_depth = None ) : <EOL> return self . root . find_bfs ( val , max_depth = max_depth ) <EOL> def find_all_bfs ( self , val , max_depth = None ) : <EOL> return self . root . find_all_bfs ( val , max_depth = max_depth ) <EOL> def dump ( self , prefix = '<STR_LIT>' , max_depth = None ) : <EOL> return self . root . dump ( prefix = prefix , max_depth = max_depth ) <EOL> def __str__ ( self ) : <EOL> return self . format ( ) <EOL> def format ( self , max_depth = None ) : <EOL> return self . root . format ( max_depth = max_depth ) <EOL> class STNode ( object ) : <EOL> OPERATORS = set ( [ token . PLUS , token . MINUS , token . STAR , token . DOUBLESTAR , <EOL> token . SLASH , token . DOUBLESLASH , token . PERCENT , <EOL> token . LEFTSHIFT , token . RIGHTSHIFT , token . AMPER , token . VBAR , <EOL> token . CIRCUMFLEX , token . TILDE , token . LESS , token . GREATER , <EOL> token . LESSEQUAL , token . GREATEREQUAL , token . EQEQUAL , <EOL> token . NOTEQUAL , token . RARROW ] ) <EOL> def __init__ ( self , tok , parent = None ) : <EOL> self . toknum = int ( tok [ <NUM_LIT:0> ] ) <EOL> self . parent = parent <EOL> self . value = None <EOL> self . children = [ ] <EOL> if self . toknum in token . tok_name : <EOL> self . name = token . tok_name [ self . toknum ] <EOL> self . value = tok [ <NUM_LIT:1> ] <EOL> if len ( tok ) > <NUM_LIT:2> : <EOL> raise ValueError ( "<STR_LIT>" . format ( tok ) ) <EOL> elif self . toknum in symbol . sym_name : <EOL> self . name = symbol . sym_name [ self . toknum ] <EOL> for child_tok in tok [ <NUM_LIT:1> : ] : <EOL> self . children . append ( STNode ( child_tok , parent = self ) ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" , tok ) <EOL> @ property <EOL> def isleaf ( self ) : <EOL> return len ( self . children ) == <NUM_LIT:0> <EOL> @ property <EOL> def isroot ( self ) : <EOL> return self . parent is None <EOL> @ property <EOL> def siblings ( self ) : <EOL> if self . isroot : <EOL> return [ ] <EOL> else : <EOL> return self . parent . children <EOL> @ property <EOL> def index ( self ) : <EOL> for idx , sib in enumerate ( self . siblings ) : <EOL> if sib is self : <EOL> return idx <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> @ property <EOL> def idx ( self ) : <EOL> return self . index <EOL> @ staticmethod <EOL> def _get_valid_maxdepth ( max_depth ) : <EOL> if max_depth is None : <EOL> max_depth = sys . maxsize <EOL> return max_depth <EOL> @ staticmethod <EOL> def default_condition ( node , val ) : <EOL> if isinstance ( val , ( list , tuple ) ) : <EOL> return node . toknum in val or node . name in val <EOL> else : <EOL> return node . toknum == val or node . name == val <EOL> def _get_valid_condition ( self , condition ) : <EOL> if condition is None : <EOL> condition = self . default_condition <EOL> return condition <EOL> def find ( self , val , max_depth = None , condition = None ) : <EOL> """<STR_LIT>""" <EOL> l = self . find_all ( val , max_depth = max_depth , condition = condition , <EOL> _saf = True ) <EOL> if l : <EOL> return l [ <NUM_LIT:0> ] <EOL> else : <EOL> return None <EOL> def find_bfs ( self , val , max_depth = None , condition = None ) : <EOL> """<STR_LIT>""" <EOL> l = self . find_all_bfs ( val , max_depth = max_depth , condition = condition , <EOL> _saf = True ) <EOL> if l : <EOL> return l [ <NUM_LIT:0> ] <EOL> else : <EOL> return None <EOL> def find_all ( self , val , max_depth = None , condition = None , _saf = False ) : <EOL> """<STR_LIT>""" <EOL> max_depth = self . _get_valid_maxdepth ( max_depth ) <EOL> condition = self . _get_valid_condition ( condition ) <EOL> if max_depth < <NUM_LIT:0> : <EOL> return [ ] <EOL> ret = [ ] <EOL> if condition ( self , val ) : <EOL> ret . append ( self ) <EOL> if _saf and ret : <EOL> return ret <EOL> for child in self . children : <EOL> ret . extend ( child . find_all ( val , max_depth = max_depth - <NUM_LIT:1> , <EOL> condition = condition , _saf = _saf ) ) <EOL> if _saf and ret : <EOL> return ret <EOL> return ret <EOL> def find_all_bfs ( self , val , max_depth = None , condition = None , _saf = False ) : <EOL> """<STR_LIT>""" <EOL> max_depth = self . _get_valid_maxdepth ( max_depth ) <EOL> condition = self . _get_valid_condition ( condition ) <EOL> ret = [ ] <EOL> depth = <NUM_LIT:0> <EOL> this_level = [ self ] <EOL> while this_level and depth <= max_depth : <EOL> ret . extend ( [ n for n in this_level if condition ( n , val ) ] ) <EOL> if _saf and ret : <EOL> break <EOL> next_level = [ ] <EOL> for n in this_level : <EOL> next_level . extend ( n . children ) <EOL> this_level = next_level <EOL> depth += <NUM_LIT:1> <EOL> return ret <EOL> def find_parent ( self , val , max_depth = None , condition = None ) : <EOL> max_depth = self . _get_valid_maxdepth ( max_depth ) <EOL> condition = self . _get_valid_condition ( condition ) <EOL> if condition ( self , val ) : <EOL> return self <EOL> elif self . isroot : <EOL> return None <EOL> else : <EOL> return self . parent . find_parent ( val , max_depth = max_depth , <EOL> condition = condition ) <EOL> def __str__ ( self ) : <EOL> s = "<STR_LIT>" <EOL> if self . isleaf : <EOL> if self . toknum == token . COMMA : <EOL> s = "<STR_LIT:U+002CU+0020>" <EOL> elif self . toknum == token . COLON : <EOL> if self . find_parent ( '<STR_LIT>' ) : <EOL> s = "<STR_LIT::>" <EOL> else : <EOL> s = "<STR_LIT>" <EOL> elif self . toknum == token . EQUAL : <EOL> if self . find_parent ( '<STR_LIT>' ) : <EOL> s = "<STR_LIT:=>" <EOL> else : <EOL> s = "<STR_LIT>" <EOL> elif self . toknum in ( token . STAR , token . DOUBLESTAR ) : <EOL> if self . find_parent ( '<STR_LIT>' , max_depth = <NUM_LIT:1> ) : <EOL> s = "<STR_LIT>" . format ( self . value ) <EOL> else : <EOL> s = "<STR_LIT>" . format ( self . value ) <EOL> elif self . toknum in self . OPERATORS : <EOL> s = "<STR_LIT>" . format ( self . value ) <EOL> elif self . toknum == token . NAME : <EOL> s = "<STR_LIT>" . format ( self . value ) <EOL> else : <EOL> s = "<STR_LIT>" . format ( self . value ) <EOL> else : <EOL> s = self . format ( ) <EOL> return s <EOL> def format ( self , max_depth = None , start = None , stop = None , ends = True ) : <EOL> max_depth = self . _get_valid_maxdepth ( max_depth ) <EOL> condition = lambda node , _ : node . isleaf <EOL> leaves = self . find_all ( <NUM_LIT:0> , max_depth = max_depth , condition = condition ) <EOL> if start : <EOL> for i , leaf in enumerate ( leaves ) : <EOL> if leaf is start : <EOL> if not ends : <EOL> i += <NUM_LIT:1> <EOL> leaves = leaves [ i : ] <EOL> break <EOL> if stop : <EOL> for i , leaf in enumerate ( leaves ) : <EOL> if leaf is stop : <EOL> if ends : <EOL> i += <NUM_LIT:1> <EOL> leaves = leaves [ : i ] <EOL> break <EOL> for i in range ( len ( leaves ) - <NUM_LIT:1> , <NUM_LIT:0> , - <NUM_LIT:1> ) : <EOL> if leaves [ i ] . toknum == leaves [ i - <NUM_LIT:1> ] . toknum == token . NAME : <EOL> leaves . insert ( i , "<STR_LIT:U+0020>" ) <EOL> return "<STR_LIT>" . join ( str ( n ) for n in leaves ) <EOL> def dump ( self , prefix = '<STR_LIT>' , max_depth = None ) : <EOL> max_depth = self . _get_valid_maxdepth ( max_depth ) <EOL> if self . isleaf : <EOL> print ( "<STR_LIT>" . format ( prefix , token . tok_name [ self . toknum ] , <EOL> self . toknum , self . value ) ) <EOL> else : <EOL> print ( "<STR_LIT>" . format ( prefix , symbol . sym_name [ self . toknum ] , <EOL> self . toknum ) ) <EOL> for child in self . children : <EOL> child . dump ( prefix = prefix + "<STR_LIT>" , max_depth = max_depth - <NUM_LIT:1> ) <EOL> def get_next_sibling ( self , offset = <NUM_LIT:1> ) : <EOL> if self . isroot : <EOL> return None <EOL> try : <EOL> return self . siblings [ self . index + offset ] <EOL> except IndexError : <EOL> return None <EOL> def find_next_sibling ( self , val , condition = None , max_distance = None , <EOL> direction = <NUM_LIT:1> ) : <EOL> condition = self . _get_valid_condition ( condition ) <EOL> max_distance = self . _get_valid_maxdepth ( max_distance ) <EOL> idx = self . index <EOL> if direction > <NUM_LIT:0> : <EOL> siblings = self . siblings [ idx : idx + <NUM_LIT:1> + max_distance ] [ <NUM_LIT:1> : ] <EOL> else : <EOL> siblings = self . siblings [ idx : idx - <NUM_LIT:1> - max_distance : - <NUM_LIT:1> ] [ <NUM_LIT:1> : ] <EOL> for sib in siblings : <EOL> if condition ( sib , val ) : <EOL> return sib <EOL> return None <EOL> def _extract_type ( s , default = None ) : <EOL> try : <EOL> val = ast . literal_eval ( s ) <EOL> ret = val . __class__ . __name__ <EOL> except Exception : <EOL> ret = default <EOL> if ret in ( '<STR_LIT>' , '<STR_LIT>' ) : <EOL> ret = ret [ : - len ( '<STR_LIT>' ) ] <EOL> return ret <EOL> def _trim_enclosing ( s , quotes = True , sequence_markers = True ) : <EOL> s = s . strip ( ) <EOL> if quotes and s [ : <NUM_LIT:3> ] == s [ - <NUM_LIT:3> : ] and s [ : <NUM_LIT:3> ] in ( '<STR_LIT>' , "<STR_LIT>" ) : <EOL> s = s [ <NUM_LIT:3> : - <NUM_LIT:3> ] <EOL> elif quotes and s [ : <NUM_LIT:1> ] == s [ - <NUM_LIT:1> : ] and s [ : <NUM_LIT:1> ] in ( '<STR_LIT:">' , "<STR_LIT:'>" ) : <EOL> s = s [ <NUM_LIT:1> : - <NUM_LIT:1> ] <EOL> elif sequence_markers and s [ : <NUM_LIT:1> ] == '<STR_LIT:(>' and s [ - <NUM_LIT:1> : ] == '<STR_LIT:)>' : <EOL> s = s [ <NUM_LIT:1> : - <NUM_LIT:1> ] <EOL> elif sequence_markers and s [ : <NUM_LIT:1> ] == '<STR_LIT:[>' and s [ - <NUM_LIT:1> : ] == '<STR_LIT:]>' : <EOL> s = s [ <NUM_LIT:1> : - <NUM_LIT:1> ] <EOL> return s <EOL> def parse_funcdef ( s , trim_string_markers = True , trim_sequence_markers = True ) : <EOL> """<STR_LIT>""" <EOL> tree = STTree ( s ) <EOL> funcdef = tree . find ( '<STR_LIT>' ) <EOL> funcname = str ( funcdef . find_all ( token . NAME , max_depth = <NUM_LIT:1> ) [ <NUM_LIT:1> ] ) <EOL> params = [ ] <EOL> argslist = funcdef . find ( '<STR_LIT>' ) <EOL> if argslist : <EOL> tfpdefs = argslist . find_all ( '<STR_LIT>' , max_depth = <NUM_LIT:1> ) <EOL> else : <EOL> tfpdefs = [ ] <EOL> lone_star = argslist . find ( token . STAR , max_depth = <NUM_LIT:1> ) <EOL> lone_star_idx = lone_star . index if lone_star else sys . maxsize <EOL> for p in tfpdefs : <EOL> default_value = None <EOL> default_type = None <EOL> annotation = "<STR_LIT>" <EOL> is_vararg = False <EOL> is_kwarg = False <EOL> is_optional = False <EOL> name_node = p . find ( token . NAME , max_depth = <NUM_LIT:1> ) <EOL> name = str ( name_node ) <EOL> marker = p . find_next_sibling ( ( token . STAR , token . DOUBLESTAR ) , <EOL> max_distance = <NUM_LIT:1> , direction = - <NUM_LIT:1> ) <EOL> if marker : <EOL> if marker . toknum == token . STAR : <EOL> is_vararg = True <EOL> elif marker . toknum == token . DOUBLESTAR : <EOL> is_kwarg = True <EOL> name = str ( marker ) + name <EOL> default_marker = p . get_next_sibling ( ) <EOL> if default_marker and default_marker . toknum == token . EQUAL : <EOL> is_optional = True <EOL> default_value = default_marker . get_next_sibling ( ) . format ( ) <EOL> default_type = _extract_type ( default_value ) <EOL> annotation_token = p . find ( token . COLON ) <EOL> if annotation_token : <EOL> annotation = p . format ( start = annotation_token , ends = False ) <EOL> params . append ( dict ( name = name , default_value = default_value , <EOL> default_type = default_type , annotation = annotation , <EOL> is_optional = is_optional , is_vararg = is_vararg , <EOL> is_kwarg = is_kwarg , <EOL> kwonly = p . index > lone_star_idx + <NUM_LIT:1> ) ) <EOL> rarrow = funcdef . find ( token . RARROW , max_depth = <NUM_LIT:1> ) <EOL> if rarrow : <EOL> annot_end = rarrow . find_next_sibling ( token . COLON ) <EOL> ret_annotation = funcdef . format ( start = rarrow , stop = annot_end , <EOL> ends = False ) <EOL> else : <EOL> ret_annotation = "<STR_LIT>" <EOL> _sanitize = lambda s : _trim_enclosing ( s , quotes = trim_string_markers , <EOL> sequence_markers = trim_sequence_markers ) <EOL> if trim_string_markers : <EOL> for p in params : <EOL> if p [ '<STR_LIT>' ] : <EOL> p [ '<STR_LIT>' ] = _sanitize ( p [ '<STR_LIT>' ] ) <EOL> ret_annotation = _sanitize ( ret_annotation ) <EOL> return funcname , params , ret_annotation <EOL> def parse_classdef ( s ) : <EOL> """<STR_LIT>""" <EOL> tree = STTree ( s ) <EOL> classdef = tree . find ( '<STR_LIT>' ) <EOL> classname = str ( classdef . find_all ( token . NAME , max_depth = <NUM_LIT:1> ) [ <NUM_LIT:1> ] ) <EOL> arglist = classdef . find ( '<STR_LIT>' ) <EOL> if arglist : <EOL> base_classes = [ str ( c ) for c in arglist . find_all ( '<STR_LIT>' ) ] <EOL> else : <EOL> base_classes = [ ] <EOL> return classname , base_classes <EOL> def _main ( ) : <EOL> def test_func ( s ) : <EOL> funcname , params , ret_annotation = parse_funcdef ( s ) <EOL> print ( "<STR_LIT>" , s ) <EOL> print ( ) <EOL> print ( "<STR_LIT>" , funcname ) <EOL> print ( "<STR_LIT>" ) <EOL> for p in params : <EOL> print ( "<STR_LIT>" , p ) <EOL> print ( "<STR_LIT>" , ret_annotation ) <EOL> test_func ( """<STR_LIT>""" ) <EOL> test_func ( """<STR_LIT>""" ) <EOL> test_func ( """<STR_LIT>""" ) <EOL> def test_class ( s ) : <EOL> classname , base_classes = parse_classdef ( s ) <EOL> print ( "<STR_LIT>" , classname ) <EOL> print ( "<STR_LIT>" ) <EOL> for bc in base_classes : <EOL> print ( "<STR_LIT>" , bc ) <EOL> test_class ( """<STR_LIT>""" ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> _main ( ) </s>
