codingutf8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license utility that sorts the imports in the custom inits of transformers transformers uses init files that delay the import of an object to when it s actually needed this is to avoid the main init importing all models which would make the line import transformers very slow when the user has all optional dependencies installed the inits with delayed imports have two halves one definining a dictionary importstructure which maps modules to the name of the objects in each module and one in typechecking which looks like a normal init for typecheckers isort or ruff properly sort the second half which looks like traditionl imports the goal of this script is to sort the first half use from the root of the repo with bash python utilscustominitisort py which will autosort the imports used in make style for a check only as used in make quality run bash python utilscustominitisort py checkonly path is defined with the intent you should run this script from the root of the repo pattern that looks at the indentation in a line pattern that matches key and puts key in group 0 pattern that matches importstructurekey and puts key in group 0 pattern that matches key and puts key in group 0 pattern that matches any stuff and puts stuff in group 0 returns the indent in given line as string search reindent searchline return if search is none else search groups0 def splitcodeinindentedblocks code str indentlevel str startprompt optionalstr none endprompt optionalstr none liststr let s split the code into lines and move to startindex index 0 lines code splitn if startprompt is not none while not linesindex startswithstartprompt index 1 blocks n joinlines index else blocks this variable contains the block treated at a given time currentblock linesindex index 1 we split into blocks until we get to the endprompt or the end of the file while index lenlines and endprompt is none or not linesindex startswithendprompt we have a nonempty line with the proper indent start of a new block if lenlinesindex 0 and getindentlinesindex indentlevel store the current block in the result and rest there are two cases the line is part of the block like a closing parenthesis or not if lencurrentblock 0 and getindentcurrentblock1 startswithindentlevel line is part of the current block currentblock appendlinesindex blocks appendn joincurrentblock if index lenlines 1 currentblock linesindex 1 index 1 else currentblock else line is not part of the current block blocks appendn joincurrentblock currentblock linesindex else just add the line to the current block currentblock appendlinesindex index 1 adds current block if it s nonempty if lencurrentblock 0 blocks appendn joincurrentblock add final block after endprompt if provided if endprompt is not none and index lenlines blocks appendn joinlinesindex return blocks def ignoreunderscoreandlowercasekey callableany str callableany str def innerx return keyx lower replace return inner def sortobjectsobjects listany key optionalcallableany str none listany if no key is provided we use a noop def noopx return x if key is none key noop constants are all uppercase they go first constants obj for obj in objects if keyobj isupper classes are not all uppercase but start with a capital they go second classes obj for obj in objects if keyobj0 isupper and not keyobj isupper functions begin with a lowercase they go last functions obj for obj in objects if not keyobj0 isupper then we sort each group key1 ignoreunderscoreandlowercasekey return sortedconstants keykey1 sortedclasses keykey1 sortedfunctions keykey1 def sortobjectsinimportimportstatement str str this inner function sort imports between def replacematch imports match groups0 if there is one import only nothing to do if not in imports return fimports keys part strip replace for part in imports split we will have a final empty element if the line finished with a comma if lenkeys1 0 keys keys 1 return joinf k for k in sortobjectskeys lines importstatement splitn if lenlines 3 here we have to sort internal imports that are on several lines one per name key object1 object2 we may have to ignore one or two lines on each side idx 2 if lines1 strip else 1 keystosort i restripline searchline groups0 for i line in enumeratelinesidx idx sortedindices sortobjectskeystosort keylambda x x1 sortedlines linesx0 idx for x in sortedindices return n joinlines idx sortedlines linesidx elif lenlines 3 here we have to sort internal imports that are on one separate line key object1 object2 if rebracketcontent searchlines1 is not none lines1 rebracketcontent subreplace lines1 else keys part strip replace for part in lines1 split we will have a final empty element if the line finished with a comma if lenkeys1 0 keys keys 1 lines1 getindentlines1 joinf k for k in sortobjectskeys return n joinlines else finally we have to deal with imports fitting on one line importstatement rebracketcontent subreplace importstatement return importstatement def sortimportsfile str checkonly bool true with openfile encodingutf8 as f code f read if the file is not a custom init there is nothing to do if importstructure not in code return blocks of indent level 0 mainblocks splitcodeinindentedblocks code startpromptimportstructure endpromptif typechecking we ignore block 0 everything untils startprompt and the last block everything after endprompt for blockidx in range1 lenmainblocks 1 check if the block contains some importstructures thingy to sort block mainblocksblockidx blocklines block splitn get to the start of the imports lineidx 0 while lineidx lenblocklines and importstructure not in blocklineslineidx skip dummy import blocks if import dummy in blocklineslineidx lineidx lenblocklines else lineidx 1 if lineidx lenblocklines continue ignore beginning and last line they don t contain anything internalblockcode n joinblocklineslineidx 1 indent getindentblocklines1 slit the internal block into blocks of indent level 1 internalblocks splitcodeinindentedblocksinternalblockcode indentlevelindent we have two categories of import key list or importstructurekey appendextend pattern redirectkey if importstructure in blocklines0 else reindirectkey grab the keys but there is a trap some lines are empty or just comments keys pattern searchb groups0 if pattern searchb is not none else none for b in internalblocks we only sort the lines with a key keystosort i key for i key in enumeratekeys if key is not none sortedindices x0 for x in sortedkeystosort keylambda x x1 we reorder the blocks by leaving empty linescomments as they were and reorder the rest count 0 reorderdedblocks for i in rangeleninternalblocks if keysi is none reorderdedblocks appendinternalblocksi else block sortobjectsinimportinternalblockssortedindicescount reorderdedblocks appendblock count 1 and we put our main block back together with its first and last line mainblocksblockidx n joinblocklines lineidx reorderdedblocks blocklines1 if code n joinmainblocks if checkonly return true else printfoverwriting file with openfile w encodingutf8 as f f writen joinmainblocks def sortimportsinallinitscheckonlytrue failures for root files in os walkpathtotransformers if init py in files result sortimportsos path joinroot init py checkonlycheckonly if result failures os path joinroot init py if lenfailures 0 raise valueerrorfwould overwrite lenfailures files run make style if name main parser argparse argumentparser parser addargumentcheckonly actionstoretrue helpwhether to only check or fix style args parser parseargs sortimportsinallinitscheckonlyargs checkonly coding utf 8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license utility that sorts the imports in the custom inits of transformers transformers uses init files that delay the import of an object to when it s actually needed this is to avoid the main init importing all models which would make the line import transformers very slow when the user has all optional dependencies installed the inits with delayed imports have two halves one definining a dictionary _import_structure which maps modules to the name of the objects in each module and one in type_checking which looks like a normal init for type checkers isort or ruff properly sort the second half which looks like traditionl imports the goal of this script is to sort the first half use from the root of the repo with bash python utils custom_init_isort py which will auto sort the imports used in make style for a check only as used in make quality run bash python utils custom_init_isort py check_only path is defined with the intent you should run this script from the root of the repo pattern that looks at the indentation in a line pattern that matches key and puts key in group 0 pattern that matches _import_structure key and puts key in group 0 pattern that matches key and puts key in group 0 pattern that matches any stuff and puts stuff in group 0 returns the indent in given line as string split some code into its indented blocks starting at a given level args code str the code to split indent_level str the indent level as string to use for identifying the blocks to split start_prompt str optional if provided only starts splitting at the line where this text is end_prompt str optional if provided stops splitting at a line where this text is warning the text before start_prompt or after end_prompt if provided is not ignored just not split the input code can thus be retrieved by joining the result returns list str the list of blocks let s split the code into lines and move to start_index this variable contains the block treated at a given time we split into blocks until we get to the end_prompt or the end of the file we have a non empty line with the proper indent start of a new block store the current block in the result and rest there are two cases the line is part of the block like a closing parenthesis or not line is part of the current block line is not part of the current block just add the line to the current block adds current block if it s nonempty add final block after end_prompt if provided wraps a key function as used in a sort to lowercase and ignore underscores sort a list of objects following the rules of isort all uppercased first camel cased second and lower cased last args objects list any the list of objects to sort key callable any str optional a function taking an object as input and returning a string used to sort them by alphabetical order if not provided will default to noop so a key must be provided if the objects are not of type string returns list any the sorted list with the same elements as in the inputs if no key is provided we use a noop constants are all uppercase they go first classes are not all uppercase but start with a capital they go second functions begin with a lowercase they go last then we sort each group sorts the imports in a single import statement args import_statement str the import statement in which to sort the imports returns str the same as the input but with objects properly sorted this inner function sort imports between if there is one import only nothing to do we will have a final empty element if the line finished with a comma here we have to sort internal imports that are on several lines one per name key object1 object2 we may have to ignore one or two lines on each side here we have to sort internal imports that are on one separate line key object1 object2 we will have a final empty element if the line finished with a comma finally we have to deal with imports fitting on one line sort the imports defined in the _import_structure of a given init args file str the path to the init to check fix check_only bool optional defaults to true whether or not to just check and not auto fix the init if the file is not a custom init there is nothing to do blocks of indent level 0 we ignore block 0 everything untils start_prompt and the last block everything after end_prompt check if the block contains some _import_structure s thingy to sort get to the start of the imports skip dummy import blocks ignore beginning and last line they don t contain anything slit the internal block into blocks of indent level 1 we have two categories of import key list or _import_structure key append extend grab the keys but there is a trap some lines are empty or just comments we only sort the lines with a key we reorder the blocks by leaving empty lines comments as they were and reorder the rest and we put our main block back together with its first and last line sort the imports defined in the _import_structure of all inits in the repo args check_only bool optional defaults to true whether or not to just check and not auto fix the init
import argparse
import os
import re
from typing import Any, Callable, List, Optional



PATH_TO_TRANSFORMERS = "src/transformers"


_re_indent = re.compile(r"^(\s*)\S")

_re_direct_key = re.compile(r'^\s*"([^"]+)":')

_re_indirect_key = re.compile(r'^\s*_import_structure\["([^"]+)"\]')

_re_strip_line = re.compile(r'^\s*"([^"]+)",\s*$')

_re_bracket_content = re.compile(r"\[([^\]]+)\]")


def get_indent(line: str) -> str:
    
    search = _re_indent.search(line)
    return "" if search is None else search.groups()[0]


def split_code_in_indented_blocks(
    code: str, indent_level: str = "", start_prompt: Optional[str] = None, end_prompt: Optional[str] = None
) -> List[str]:
    
    
    index = 0
    lines = code.split("\n")
    if start_prompt is not None:
        while not lines[index].startswith(start_prompt):
            index += 1
        blocks = ["\n".join(lines[:index])]
    else:
        blocks = []

    
    current_block = [lines[index]]
    index += 1
    
    while index < len(lines) and (end_prompt is None or not lines[index].startswith(end_prompt)):
        
        if len(lines[index]) > 0 and get_indent(lines[index]) == indent_level:
            
            
            if len(current_block) > 0 and get_indent(current_block[-1]).startswith(indent_level + " "):
                
                current_block.append(lines[index])
                blocks.append("\n".join(current_block))
                if index < len(lines) - 1:
                    current_block = [lines[index + 1]]
                    index += 1
                else:
                    current_block = []
            else:
                
                blocks.append("\n".join(current_block))
                current_block = [lines[index]]
        else:
            
            current_block.append(lines[index])
        index += 1

    
    if len(current_block) > 0:
        blocks.append("\n".join(current_block))

    
    if end_prompt is not None and index < len(lines):
        blocks.append("\n".join(lines[index:]))

    return blocks


def ignore_underscore_and_lowercase(key: Callable[[Any], str]) -> Callable[[Any], str]:
    

    def _inner(x):
        return key(x).lower().replace("_", "")

    return _inner


def sort_objects(objects: List[Any], key: Optional[Callable[[Any], str]] = None) -> List[Any]:
    

    
    def noop(x):
        return x

    if key is None:
        key = noop
    
    constants = [obj for obj in objects if key(obj).isupper()]
    
    classes = [obj for obj in objects if key(obj)[0].isupper() and not key(obj).isupper()]
    
    functions = [obj for obj in objects if not key(obj)[0].isupper()]

    
    key1 = ignore_underscore_and_lowercase(key)
    return sorted(constants, key=key1) + sorted(classes, key=key1) + sorted(functions, key=key1)


def sort_objects_in_import(import_statement: str) -> str:
    

    
    def _replace(match):
        imports = match.groups()[0]
        
        if "," not in imports:
            return f"[{imports}]"
        keys = [part.strip().replace('"', "") for part in imports.split(",")]
        
        if len(keys[-1]) == 0:
            keys = keys[:-1]
        return "[" + ", ".join([f'"{k}"' for k in sort_objects(keys)]) + "]"

    lines = import_statement.split("\n")
    if len(lines) > 3:
        
        
        
        
        
        

        
        idx = 2 if lines[1].strip() == "[" else 1
        keys_to_sort = [(i, _re_strip_line.search(line).groups()[0]) for i, line in enumerate(lines[idx:-idx])]
        sorted_indices = sort_objects(keys_to_sort, key=lambda x: x[1])
        sorted_lines = [lines[x[0] + idx] for x in sorted_indices]
        return "\n".join(lines[:idx] + sorted_lines + lines[-idx:])
    elif len(lines) == 3:
        
        
        
        
        if _re_bracket_content.search(lines[1]) is not None:
            lines[1] = _re_bracket_content.sub(_replace, lines[1])
        else:
            keys = [part.strip().replace('"', "") for part in lines[1].split(",")]
            
            if len(keys[-1]) == 0:
                keys = keys[:-1]
            lines[1] = get_indent(lines[1]) + ", ".join([f'"{k}"' for k in sort_objects(keys)])
        return "\n".join(lines)
    else:
        
        import_statement = _re_bracket_content.sub(_replace, import_statement)
        return import_statement


def sort_imports(file: str, check_only: bool = True):
    
    with open(file, encoding="utf-8") as f:
        code = f.read()

    
    if "_import_structure" not in code:
        return

    
    main_blocks = split_code_in_indented_blocks(
        code, start_prompt="_import_structure = {", end_prompt="if TYPE_CHECKING:"
    )

    
    for block_idx in range(1, len(main_blocks) - 1):
        
        block = main_blocks[block_idx]
        block_lines = block.split("\n")

        
        line_idx = 0
        while line_idx < len(block_lines) and "_import_structure" not in block_lines[line_idx]:
            
            if "import dummy" in block_lines[line_idx]:
                line_idx = len(block_lines)
            else:
                line_idx += 1
        if line_idx >= len(block_lines):
            continue

        
        internal_block_code = "\n".join(block_lines[line_idx:-1])
        indent = get_indent(block_lines[1])
        
        internal_blocks = split_code_in_indented_blocks(internal_block_code, indent_level=indent)
        
        pattern = _re_direct_key if "_import_structure = {" in block_lines[0] else _re_indirect_key
        
        keys = [(pattern.search(b).groups()[0] if pattern.search(b) is not None else None) for b in internal_blocks]
        
        keys_to_sort = [(i, key) for i, key in enumerate(keys) if key is not None]
        sorted_indices = [x[0] for x in sorted(keys_to_sort, key=lambda x: x[1])]

        
        count = 0
        reorderded_blocks = []
        for i in range(len(internal_blocks)):
            if keys[i] is None:
                reorderded_blocks.append(internal_blocks[i])
            else:
                block = sort_objects_in_import(internal_blocks[sorted_indices[count]])
                reorderded_blocks.append(block)
                count += 1

        
        main_blocks[block_idx] = "\n".join(block_lines[:line_idx] + reorderded_blocks + [block_lines[-1]])

    if code != "\n".join(main_blocks):
        if check_only:
            return True
        else:
            print(f"Overwriting {file}.")
            with open(file, "w", encoding="utf-8") as f:
                f.write("\n".join(main_blocks))


def sort_imports_in_all_inits(check_only=True):
    
    failures = []
    for root, _, files in os.walk(PATH_TO_TRANSFORMERS):
        if "__init__.py" in files:
            result = sort_imports(os.path.join(root, "__init__.py"), check_only=check_only)
            if result:
                failures = [os.path.join(root, "__init__.py")]
    if len(failures) > 0:
        raise ValueError(f"Would overwrite {len(failures)} files, run `make style`.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--check_only", action="store_true", help="Whether to only check or fix style.")
    args = parser.parse_args()

    sort_imports_in_all_inits(check_only=args.check_only)