import { CancellationToken } from 'vscode-languageserver' ; import { compareImportStatements , getImportGroup , getTopLevelImports , ImportStatement , } from '../analyzer/importStatementUtils' ; import { throwIfCancellationRequested } from '../common/cancellationUtils' ; import { TextEditAction } from '../common/editAction' ; import { convertOffsetToPosition } from '../common/positionUtils' ; import { Range } from '../common/textRange' ; import { TextRange } from '../common/textRange' ; import { ImportAsNode , ImportFromAsNode , ImportFromNode , ParseNodeType } from '../parser/parseNodes' ; import { ParseResults } from '../parser/parser' ; const _maxLineLength = 88 ; export class ImportSorter { constructor ( private _parseResults , private _cancellationToken ) { } sort ( ) { throwIfCancellationRequested ( this . _cancellationToken ) ; const actions = [ ] ; const importStatements = getTopLevelImports ( this . _parseResults . parseTree ) ; const sortedStatements = importStatements . orderedImports . map ( ( s ) => s ) . sort ( ( a , b ) => { return compareImportStatements ( a , b ) ; } ) ; if ( sortedStatements . length === 0 ) { return [ ] ; } const primaryRange = this . _getPrimaryReplacementRange ( importStatements . orderedImports ) ; actions . push ( { range : primaryRange , replacementText : this . _generateSortedImportText ( sortedStatements ) , } ) ; this . _addSecondaryReplacementRanges ( importStatements . orderedImports , actions ) ; return actions ; } private _getPrimaryReplacementRange ( statements ) { let statementLimit = statements . findIndex ( ( s ) => s . followsNonImportStatement ) ; if ( statementLimit < 0 ) { statementLimit = statements . length ; } const lastStatement = statements [ statementLimit - 1 ] . node ; return { start : convertOffsetToPosition ( statements [ 0 ] . node . start , this . _parseResults . tokenizerOutput . lines ) , end : convertOffsetToPosition ( TextRange . getEnd ( lastStatement ) , this . _parseResults . tokenizerOutput . lines ) , } ; } private _addSecondaryReplacementRanges ( statements , actions ) { let secondaryBlockStart = statements . findIndex ( ( s ) => s . followsNonImportStatement ) ; if ( secondaryBlockStart < 0 ) { return ; } while ( true ) { let secondaryBlockLimit = statements . findIndex ( ( s , index ) => index > secondaryBlockStart && s . followsNonImportStatement ) ; if ( secondaryBlockLimit < 0 ) { secondaryBlockLimit = statements . length ; } actions . push ( { range : { start : convertOffsetToPosition ( statements [ secondaryBlockStart ] . node . start , this . _parseResults . tokenizerOutput . lines ) , end : convertOffsetToPosition ( TextRange . getEnd ( statements [ secondaryBlockLimit - 1 ] . node ) , this . _parseResults . tokenizerOutput . lines ) , } , replacementText : '' , } ) ; secondaryBlockStart = secondaryBlockLimit ; if ( secondaryBlockStart >= statements . length ) { break ; } } } private _generateSortedImportText ( sortedStatements ) { let importText = '' ; let prevImportGroup = getImportGroup ( sortedStatements [ 0 ] ) ; for ( const statement of sortedStatements ) { const curImportType = getImportGroup ( statement ) ; if ( prevImportGroup !== curImportType ) { importText += this . _parseResults . tokenizerOutput . predominantEndOfLineSequence ; prevImportGroup = curImportType ; } let importLine ; if ( statement . node . nodeType === ParseNodeType . Import ) { importLine = this . _formatImportNode ( statement . subnode ! , statement . moduleName ) ; } else { importLine = this . _formatImportFromNode ( statement . node , statement . moduleName ) ; } if ( statement !== sortedStatements [ sortedStatements . length - 1 ] ) { importLine += this . _parseResults . tokenizerOutput . predominantEndOfLineSequence ; } importText += importLine ; } return importText ; } private _formatImportNode ( subnode , moduleName ) { let importText = `import ${ moduleName }` ; if ( subnode . alias ) { importText += ` as ${ subnode . alias . value }` ; } return importText ; } private _formatImportFromNode ( node , moduleName ) { const symbols = node . imports . sort ( ( a , b ) => this . _compareSymbols ( a , b ) ) . map ( ( symbol ) => { let symbolText = symbol . name . value ; if ( symbol . alias ) { symbolText += ` as ${ symbol . alias . value }` ; } return symbolText ; } ) ; let cumulativeText = `from ${ moduleName } import ` ; if ( node . isWildcardImport ) { return cumulativeText + '*' ; } const symbolText = symbols . join ( ', ' ) ; if ( cumulativeText . length + symbolText . length <= _maxLineLength ) { return cumulativeText + symbolText ; } cumulativeText += '(' + this . _parseResults . tokenizerOutput . predominantEndOfLineSequence ; for ( const symbol of symbols ) { cumulativeText += this . _parseResults . tokenizerOutput . predominantTabSequence + symbol + ',' + this . _parseResults . tokenizerOutput . predominantEndOfLineSequence ; } cumulativeText += ')' ; return cumulativeText ; } private _compareSymbols ( a , b ) { return a . name . value < b . name . value ? - 1 : 1 ; } }