<s> """<STR_LIT>""" <EOL> from __future__ import print_function <EOL> import argparse <EOL> import ast <EOL> import collections <EOL> import os <EOL> import sys <EOL> ParsedImport = collections . namedtuple ( <EOL> '<STR_LIT>' , [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> _REQUIRED_FUTURE_IMPORTS = set ( <EOL> [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> class ProcessingException ( Exception ) : <EOL> def __init__ ( self , description ) : <EOL> super ( ProcessingException , self ) . __init__ ( description ) <EOL> def main ( ) : <EOL> argparser = argparse . ArgumentParser ( ) <EOL> argparser . add_argument ( <EOL> '<STR_LIT>' , <EOL> type = str , <EOL> nargs = "<STR_LIT:+>" ) <EOL> args = argparser . parse_args ( ) <EOL> modules = sorted ( args . modules ) <EOL> for m in modules : <EOL> try : <EOL> process_module ( m ) <EOL> except Exception as e : <EOL> print ( '<STR_LIT>' , m ) <EOL> print ( e ) <EOL> return <NUM_LIT:1> <EOL> return <NUM_LIT:0> <EOL> def process_module ( m ) : <EOL> raw = '<STR_LIT>' . join ( open ( m ) ) <EOL> new_raw = rewrite_module ( raw , os . path . basename ( m ) ) <EOL> if new_raw != raw : <EOL> with open ( m , '<STR_LIT:w>' ) as f : <EOL> f . write ( new_raw ) <EOL> print ( "<STR_LIT>" , m ) <EOL> def rewrite_module ( s , module_name ) : <EOL> imports = parse_imports ( s ) <EOL> groups = group_imports ( imports , module_name ) <EOL> lines = s . splitlines ( ) <EOL> new_import_lines = None <EOL> if groups : <EOL> new_import_lines = groups [ <NUM_LIT:0> ] <EOL> for g in groups [ <NUM_LIT:1> : ] : <EOL> new_import_lines . append ( '<STR_LIT>' ) <EOL> new_import_lines += g <EOL> if new_import_lines : <EOL> if imports : <EOL> first_line = imports [ <NUM_LIT:0> ] . line - <NUM_LIT:1> <EOL> last_line = imports [ - <NUM_LIT:1> ] . line - <NUM_LIT:1> <EOL> lines [ first_line : last_line + <NUM_LIT:1> ] = new_import_lines <EOL> else : <EOL> lines [ <NUM_LIT:0> : <NUM_LIT:1> ] = lines [ <NUM_LIT:0> : <NUM_LIT:1> ] + new_import_lines <EOL> return '<STR_LIT:\n>' . join ( lines ) + '<STR_LIT:\n>' <EOL> def parse_imports ( s ) : <EOL> module = ast . parse ( s ) <EOL> imports = [ ] <EOL> for c in ast . iter_child_nodes ( module ) : <EOL> if isinstance ( c , ast . ImportFrom ) : <EOL> fields = dict ( ast . iter_fields ( c ) ) <EOL> from_module = fields [ '<STR_LIT>' ] <EOL> if from_module != '<STR_LIT>' : <EOL> raise ProcessingException ( <EOL> '<STR_LIT>' + str ( c . lineno ) ) <EOL> names = [ dict ( ast . iter_fields ( i ) ) [ '<STR_LIT:name>' ] for i in fields [ '<STR_LIT>' ] ] <EOL> if len ( names ) != <NUM_LIT:1> : <EOL> raise ProcessingException ( <EOL> '<STR_LIT>' + str ( c . lineno ) ) <EOL> imports . append ( ParsedImport ( c . lineno , names [ <NUM_LIT:0> ] , True ) ) <EOL> if isinstance ( c , ast . Import ) : <EOL> fields = dict ( ast . iter_fields ( c ) ) <EOL> names = [ dict ( ast . iter_fields ( i ) ) [ '<STR_LIT:name>' ] for i in fields [ '<STR_LIT>' ] ] <EOL> if len ( names ) != <NUM_LIT:1> : <EOL> raise ProcessingException ( <EOL> '<STR_LIT>' + str ( c . lineno ) ) <EOL> imports . append ( ParsedImport ( c . lineno , names [ <NUM_LIT:0> ] , False ) ) <EOL> return imports <EOL> def group_imports ( imports , module_name ) : <EOL> groups = [ ] <EOL> package_group = module_name [ <NUM_LIT:0> : <NUM_LIT:3> ] <EOL> package = module_name [ <NUM_LIT:0> : module_name . index ( '<STR_LIT:_>' ) ] <EOL> current_group = [ i . module for i in imports if i . is_future ] <EOL> current_group = list ( set ( current_group ) | _REQUIRED_FUTURE_IMPORTS ) <EOL> imports = [ i . module for i in imports ] <EOL> push_current_future_group ( groups , current_group , imports ) <EOL> current_group = [ i for i in imports if '<STR_LIT:_>' not in i ] <EOL> push_current_group ( groups , current_group , imports ) <EOL> current_group = [ i for i in imports if not i . startswith ( package_group ) ] <EOL> push_current_group ( groups , current_group , imports ) <EOL> current_group = [ i for i in imports if not i . startswith ( package ) ] <EOL> push_current_group ( groups , current_group , imports ) <EOL> current_group = [ i for i in imports if i . startswith ( package ) ] <EOL> push_current_group ( groups , current_group , imports ) <EOL> assert not imports <EOL> return groups <EOL> def push_current_group ( groups , current_group , imports ) : <EOL> if current_group : <EOL> imports [ : ] = list ( set ( imports ) - set ( current_group ) ) <EOL> current_group = [ '<STR_LIT>' + i for i in current_group ] <EOL> current_group . sort ( ) <EOL> groups . append ( current_group ) <EOL> def push_current_future_group ( groups , current_group , imports ) : <EOL> if current_group : <EOL> imports [ : ] = list ( set ( imports ) - set ( current_group ) ) <EOL> current_group = [ '<STR_LIT>' + i for i in current_group ] <EOL> current_group . sort ( ) <EOL> groups . append ( current_group ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> sys . exit ( main ( ) ) </s>
