import { CancellationToken } from 'vscode-languageserver' ; import { getTextEditsForAutoImportInsertion , getTextEditsForAutoImportSymbolAddition , getTopLevelImports , ImportGroup , } from '../analyzer/importStatementUtils' ; import * as ParseTreeUtils from '../analyzer/parseTreeUtils' ; import { Commands } from '../commands/commands' ; import { throwIfCancellationRequested } from '../common/cancellationUtils' ; import { TextEditAction } from '../common/editAction' ; import { convertOffsetToPosition } from '../common/positionUtils' ; import { TextRange } from '../common/textRange' ; import { ParseNode , ParseNodeType } from '../parser/parseNodes' ; import { ParseResults } from '../parser/parser' ; import { ImportSorter } from './importSorter' ; export function performQuickAction ( command , args , parseResults , token ) { if ( command === Commands . orderImports ) { const importSorter = new ImportSorter ( parseResults , token ) ; return importSorter . sort ( ) ; } else if ( command === Commands . addMissingOptionalToParam ) { if ( args . length >= 1 ) { const nodeOffset = parseInt ( args [ 0 ] , 10 ) ; return _addMissingOptionalToParam ( parseResults , nodeOffset , token ) ; } } return [ ] ; } function _addMissingOptionalToParam ( parseResults , offset , token ) { throwIfCancellationRequested ( token ) ; let node = ParseTreeUtils . findNodeByOffset ( parseResults . parseTree , offset ) ; while ( node ) { if ( node . nodeType === ParseNodeType . Parameter ) { break ; } node = node . parent ; } if ( ! node ) { return [ ] ; } const typeAnnotation = node . typeAnnotation || node . typeAnnotationComment ; if ( ! typeAnnotation ) { return [ ] ; } const editActions = [ ] ; const startPos = convertOffsetToPosition ( typeAnnotation . start , parseResults . tokenizerOutput . lines ) ; const endPos = convertOffsetToPosition ( TextRange . getEnd ( typeAnnotation ) , parseResults . tokenizerOutput . lines ) ; editActions . push ( { range : { start : startPos , end : startPos } , replacementText : 'Optional[' , } ) ; editActions . push ( { range : { start : endPos , end : endPos } , replacementText : ']' , } ) ; const importStatements = getTopLevelImports ( parseResults . parseTree ) ; const importStatement = importStatements . orderedImports . find ( ( imp ) => imp . moduleName === 'typing' ) ; if ( importStatement && importStatement . node . nodeType === ParseNodeType . ImportFrom && ! importStatement . node . isWildcardImport ) { const additionalEditActions = getTextEditsForAutoImportSymbolAddition ( { name : 'Optional' } , importStatement , parseResults ) ; editActions . push ( ... additionalEditActions ) ; } else { const additionalEditActions = getTextEditsForAutoImportInsertion ( { name : 'Optional' } , importStatements , 'typing' , ImportGroup . BuiltIn , parseResults , startPos ) ; editActions . push ( ... additionalEditActions ) ; } return editActions ; }