<s> from __future__ import absolute_import <EOL> from __future__ import print_function <EOL> from __future__ import unicode_literals <EOL> import argparse <EOL> import ast <EOL> import collections <EOL> import difflib <EOL> import functools <EOL> import io <EOL> import sys <EOL> import tokenize <EOL> import six <EOL> from aspy . refactor_imports . import_obj import import_obj_from_str <EOL> from aspy . refactor_imports . sort import sort <EOL> class CodeType ( object ) : <EOL> PRE_IMPORT_CODE = '<STR_LIT>' <EOL> IMPORT = '<STR_LIT>' <EOL> NON_CODE = '<STR_LIT>' <EOL> CODE = '<STR_LIT:code>' <EOL> CodePartition = collections . namedtuple ( '<STR_LIT>' , ( '<STR_LIT>' , '<STR_LIT:src>' ) ) <EOL> TERMINATES_COMMENT = frozenset ( ( tokenize . NL , tokenize . ENDMARKER ) ) <EOL> TERMINATES_DOCSTRING = frozenset ( ( tokenize . NEWLINE , tokenize . ENDMARKER ) ) <EOL> TERMINATES_IMPORT = frozenset ( ( tokenize . NEWLINE , tokenize . ENDMARKER ) ) <EOL> class TopLevelImportVisitor ( ast . NodeVisitor ) : <EOL> def __init__ ( self ) : <EOL> self . top_level_import_line_numbers = [ ] <EOL> def _visit_import ( self , node ) : <EOL> if node . col_offset == <NUM_LIT:0> : <EOL> self . top_level_import_line_numbers . append ( node . lineno ) <EOL> visit_Import = visit_ImportFrom = _visit_import <EOL> def get_line_offsets_by_line_no ( src ) : <EOL> offsets = [ None , <NUM_LIT:0> ] <EOL> for line in src . splitlines ( ) : <EOL> offsets . append ( offsets [ - <NUM_LIT:1> ] + len ( line ) + <NUM_LIT:1> ) <EOL> return offsets <EOL> def _partitions_to_src ( partitions ) : <EOL> return '<STR_LIT>' . join ( part . src for part in partitions ) <EOL> def partition_source ( src ) : <EOL> """<STR_LIT>""" <EOL> if type ( src ) is not six . text_type : <EOL> raise TypeError ( '<STR_LIT>' . format ( type ( src ) ) ) <EOL> ast_obj = ast . parse ( src . encode ( '<STR_LIT>' ) ) <EOL> visitor = TopLevelImportVisitor ( ) <EOL> visitor . visit ( ast_obj ) <EOL> line_offsets = get_line_offsets_by_line_no ( src ) <EOL> chunks = [ ] <EOL> startpos = <NUM_LIT:0> <EOL> pending_chunk_type = None <EOL> possible_ending_tokens = None <EOL> seen_import = False <EOL> for ( <EOL> token_type , token_text , ( srow , scol ) , ( erow , ecol ) , _ , <EOL> ) in tokenize . generate_tokens ( io . StringIO ( src ) . readline ) : <EOL> if pending_chunk_type is None : <EOL> if not seen_import and token_type == tokenize . COMMENT : <EOL> if '<STR_LIT>' in token_text : <EOL> chunks . append ( CodePartition ( CodeType . CODE , src [ startpos : ] ) ) <EOL> break <EOL> else : <EOL> pending_chunk_type = CodeType . PRE_IMPORT_CODE <EOL> possible_ending_tokens = TERMINATES_COMMENT <EOL> elif not seen_import and token_type == tokenize . STRING : <EOL> pending_chunk_type = CodeType . PRE_IMPORT_CODE <EOL> possible_ending_tokens = TERMINATES_DOCSTRING <EOL> elif scol == <NUM_LIT:0> and srow in visitor . top_level_import_line_numbers : <EOL> seen_import = True <EOL> pending_chunk_type = CodeType . IMPORT <EOL> possible_ending_tokens = TERMINATES_IMPORT <EOL> elif token_type == tokenize . NL : <EOL> endpos = line_offsets [ erow ] + ecol <EOL> srctext = src [ startpos : endpos ] <EOL> startpos = endpos <EOL> chunks . append ( CodePartition ( CodeType . NON_CODE , srctext ) ) <EOL> elif token_type == tokenize . COMMENT : <EOL> if '<STR_LIT>' in token_text : <EOL> chunks . append ( CodePartition ( CodeType . CODE , src [ startpos : ] ) ) <EOL> break <EOL> else : <EOL> pending_chunk_type = CodeType . CODE <EOL> possible_ending_tokens = TERMINATES_COMMENT <EOL> elif token_type == tokenize . ENDMARKER : <EOL> pass <EOL> else : <EOL> chunks . append ( CodePartition ( CodeType . CODE , src [ startpos : ] ) ) <EOL> break <EOL> elif token_type in possible_ending_tokens : <EOL> endpos = line_offsets [ erow ] + ecol <EOL> srctext = src [ startpos : endpos ] <EOL> startpos = endpos <EOL> chunks . append ( CodePartition ( pending_chunk_type , srctext ) ) <EOL> pending_chunk_type = None <EOL> possible_ending_tokens = None <EOL> elif token_type == tokenize . COMMENT and '<STR_LIT>' in token_text : <EOL> chunks . append ( CodePartition ( CodeType . CODE , src [ startpos : ] ) ) <EOL> break <EOL> assert _partitions_to_src ( chunks ) == src <EOL> return chunks <EOL> def combine_trailing_code_chunks ( partitions ) : <EOL> chunks = list ( partitions ) <EOL> NON_COMBINABLE = ( CodeType . IMPORT , CodeType . PRE_IMPORT_CODE ) <EOL> if chunks and chunks [ - <NUM_LIT:1> ] . code_type not in NON_COMBINABLE : <EOL> src = chunks . pop ( ) . src <EOL> while chunks and chunks [ - <NUM_LIT:1> ] . code_type not in NON_COMBINABLE : <EOL> src = chunks . pop ( ) . src + src <EOL> chunks . append ( CodePartition ( CodeType . CODE , src ) ) <EOL> return chunks <EOL> def separate_comma_imports ( partitions ) : <EOL> """<STR_LIT>""" <EOL> def _inner ( ) : <EOL> for partition in partitions : <EOL> if partition . code_type is CodeType . IMPORT : <EOL> import_obj = import_obj_from_str ( partition . src ) <EOL> if import_obj . has_multiple_imports : <EOL> for new_import_obj in import_obj . split_imports ( ) : <EOL> yield CodePartition ( <EOL> CodeType . IMPORT , new_import_obj . to_text ( ) <EOL> ) <EOL> else : <EOL> yield partition <EOL> else : <EOL> yield partition <EOL> return list ( _inner ( ) ) <EOL> def add_imports ( partitions , to_add = ( ) ) : <EOL> partitions = list ( partitions ) <EOL> if not _partitions_to_src ( partitions ) . strip ( ) : <EOL> return partitions <EOL> if not partitions [ - <NUM_LIT:1> ] . src . endswith ( '<STR_LIT:\n>' ) : <EOL> partitions [ - <NUM_LIT:1> ] = CodePartition ( <EOL> partitions [ - <NUM_LIT:1> ] . code_type , <EOL> partitions [ - <NUM_LIT:1> ] . src + '<STR_LIT:\n>' , <EOL> ) <EOL> return partitions + [ <EOL> CodePartition ( CodeType . IMPORT , imp_statement . strip ( ) + '<STR_LIT:\n>' ) <EOL> for imp_statement in to_add <EOL> ] <EOL> def remove_imports ( partitions , to_remove = ( ) ) : <EOL> to_remove_imports = set ( <EOL> import_obj_from_str ( imp_statement ) for imp_statement in to_remove <EOL> ) <EOL> def _inner ( ) : <EOL> for partition in partitions : <EOL> if ( <EOL> partition . code_type is not CodeType . IMPORT or <EOL> import_obj_from_str ( partition . src ) not in to_remove_imports <EOL> ) : <EOL> yield partition <EOL> return list ( _inner ( ) ) <EOL> def remove_duplicated_imports ( partitions ) : <EOL> def _inner ( ) : <EOL> seen = set ( ) <EOL> for partition in partitions : <EOL> if partition . code_type is CodeType . IMPORT : <EOL> import_obj = import_obj_from_str ( partition . src ) <EOL> if import_obj not in seen : <EOL> seen . add ( import_obj ) <EOL> yield partition <EOL> else : <EOL> yield partition <EOL> return list ( _inner ( ) ) <EOL> def apply_import_sorting ( partitions , ** sort_kwargs ) : <EOL> pre_import_code = [ ] <EOL> imports = [ ] <EOL> trash = [ ] <EOL> rest = [ ] <EOL> for partition in partitions : <EOL> { <EOL> CodeType . PRE_IMPORT_CODE : pre_import_code , <EOL> CodeType . IMPORT : imports , <EOL> CodeType . NON_CODE : trash , <EOL> CodeType . CODE : rest , <EOL> } [ partition . code_type ] . append ( partition ) <EOL> imports = [ <EOL> partition if partition . src . endswith ( '<STR_LIT:\n>' ) else <EOL> CodePartition ( CodeType . IMPORT , partition . src + '<STR_LIT:\n>' ) <EOL> for partition in imports <EOL> ] <EOL> import_obj_to_partition = dict ( <EOL> ( import_obj_from_str ( partition . src ) , partition ) <EOL> for partition in imports <EOL> ) <EOL> new_imports = [ ] <EOL> sorted_blocks = sort ( import_obj_to_partition . keys ( ) , ** sort_kwargs ) <EOL> for block in sorted_blocks : <EOL> for import_obj in block : <EOL> new_imports . append ( import_obj_to_partition [ import_obj ] ) <EOL> new_imports . append ( CodePartition ( CodeType . NON_CODE , '<STR_LIT:\n>' ) ) <EOL> if new_imports : <EOL> new_imports . pop ( ) <EOL> restsrc = _partitions_to_src ( rest ) <EOL> restsrc = restsrc . rstrip ( ) <EOL> if restsrc : <EOL> rest = [ <EOL> CodePartition ( CodeType . CODE , restsrc + '<STR_LIT:\n>' ) , <EOL> ] <EOL> else : <EOL> rest = [ ] <EOL> return pre_import_code + new_imports + rest <EOL> def _get_steps ( imports_to_add , imports_to_remove , ** sort_kwargs ) : <EOL> yield combine_trailing_code_chunks <EOL> yield separate_comma_imports <EOL> if imports_to_add : <EOL> yield functools . partial ( add_imports , to_add = imports_to_add ) <EOL> if imports_to_remove : <EOL> yield functools . partial ( remove_imports , to_remove = imports_to_remove ) <EOL> yield remove_duplicated_imports <EOL> yield functools . partial ( apply_import_sorting , ** sort_kwargs ) <EOL> def fix_file_contents ( <EOL> contents , imports_to_add = ( ) , imports_to_remove = ( ) , ** sort_kwargs <EOL> ) : <EOL> partitioned = partition_source ( contents ) <EOL> for step in _get_steps ( imports_to_add , imports_to_remove , ** sort_kwargs ) : <EOL> partitioned = step ( partitioned ) <EOL> return _partitions_to_src ( partitioned ) <EOL> def report_diff ( contents , new_contents , filename ) : <EOL> diff = '<STR_LIT>' . join ( difflib . unified_diff ( <EOL> io . StringIO ( contents ) . readlines ( ) , <EOL> io . StringIO ( new_contents ) . readlines ( ) , <EOL> fromfile = filename , tofile = filename , <EOL> ) ) <EOL> if not diff . endswith ( '<STR_LIT:\n>' ) : <EOL> diff += '<STR_LIT>' <EOL> print ( diff , end = '<STR_LIT>' ) <EOL> def apply_reordering ( new_contents , filename ) : <EOL> print ( '<STR_LIT>' . format ( filename ) ) <EOL> with io . open ( filename , '<STR_LIT:w>' ) as f : <EOL> f . write ( new_contents ) <EOL> def main ( argv = None ) : <EOL> parser = argparse . ArgumentParser ( ) <EOL> parser . add_argument ( '<STR_LIT>' , nargs = '<STR_LIT:*>' ) <EOL> group = parser . add_mutually_exclusive_group ( required = False ) <EOL> group . add_argument ( <EOL> '<STR_LIT>' , action = '<STR_LIT:store_true>' , <EOL> help = '<STR_LIT>' , <EOL> ) <EOL> group . add_argument ( <EOL> '<STR_LIT>' , action = '<STR_LIT:store_true>' , <EOL> help = '<STR_LIT>' , <EOL> ) <EOL> parser . add_argument ( <EOL> '<STR_LIT>' , action = '<STR_LIT>' , <EOL> help = '<STR_LIT>' , <EOL> ) <EOL> parser . add_argument ( <EOL> '<STR_LIT>' , action = '<STR_LIT>' , <EOL> help = ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> ) , <EOL> ) <EOL> parser . add_argument ( <EOL> '<STR_LIT>' , default = '<STR_LIT:.>' , <EOL> help = ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> ) , <EOL> ) <EOL> args = parser . parse_args ( argv ) <EOL> retv = <NUM_LIT:0> <EOL> for filename in args . filenames : <EOL> contents = io . open ( filename ) . read ( ) <EOL> new_contents = fix_file_contents ( <EOL> contents , <EOL> imports_to_add = args . add_import , <EOL> imports_to_remove = args . remove_import , <EOL> application_directories = args . application_directories . split ( '<STR_LIT::>' ) , <EOL> ) <EOL> if contents != new_contents : <EOL> retv = <NUM_LIT:1> <EOL> if args . diff_only : <EOL> report_diff ( contents , new_contents , filename ) <EOL> elif args . print_only : <EOL> print ( '<STR_LIT>' . format ( filename ) , file = sys . stderr ) <EOL> print ( new_contents , end = '<STR_LIT>' ) <EOL> else : <EOL> apply_reordering ( new_contents , filename ) <EOL> return retv <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> exit ( main ( ) ) </s>
