import { CancellationToken } from 'vscode-languageserver' ; import { getFileInfo } from '../analyzer/analyzerNodeInfo' ; import { Declaration , isAliasDeclaration , isClassDeclaration , isFunctionDeclaration , isParameterDeclaration , isVariableDeclaration , ModuleLoaderActions , } from '../analyzer/declaration' ; import { createSynthesizedAliasDeclaration , getNameFromDeclaration , isDefinedInFile , } from '../analyzer/declarationUtils' ; import { ImportResolver } from '../analyzer/importResolver' ; import { getRelativeModuleName , getTextEditsForAutoImportInsertions , getTextEditsForAutoImportSymbolAddition , getTopLevelImports , ImportNameInfo , ImportNameWithModuleInfo , ImportStatements , } from '../analyzer/importStatementUtils' ; import { getDottedName , getDottedNameWithGivenNodeAsLastName , isLastNameOfDottedName , } from '../analyzer/parseTreeUtils' ; import { ParseTreeWalker } from '../analyzer/parseTreeWalker' ; import { ScopeType } from '../analyzer/scope' ; import { getScopeForNode } from '../analyzer/scopeUtils' ; import { TypeEvaluator } from '../analyzer/typeEvaluatorTypes' ; import { throwIfCancellationRequested } from '../common/cancellationUtils' ; import { addIfUnique , createMapFromItems , getOrAdd , removeArrayElements } from '../common/collectionUtils' ; import { ConfigOptions } from '../common/configOptions' ; import { TextEditAction } from '../common/editAction' ; import { getDirectoryPath } from '../common/pathUtils' ; import { convertOffsetToPosition } from '../common/positionUtils' ; import { TextRange } from '../common/textRange' ; import { ModuleNameNode , NameNode , ParseNode , ParseNodeType } from '../parser/parseNodes' ; import { ParseResults } from '../parser/parser' ; export interface ImportData { containsUnreferenceableSymbols : boolean ; declarations : Map < Declaration , NameNode [ ] > ; } export class ImportAdder { constructor ( private _configOptions , private _importResolver , private _evaluator ) { } collectImportsForSymbolsUsed ( parseResults , range , token ) { const collector = new NameCollector ( this . _evaluator , parseResults , range , token ) ; collector . walk ( parseResults . parseTree ) ; return { containsUnreferenceableSymbols : collector . containsUnreferenceableSymbols , declarations : collector . declsForSymbols , } ; } applyImports ( result , parseResults , insertionPosition , token ) { throwIfCancellationRequested ( token ) ; const filePath = getFileInfo ( parseResults . parseTree ) . filePath ; const importStatements = getTopLevelImports ( parseResults . parseTree ) ; const execEnv = this . _configOptions . findExecEnvironment ( filePath ) ; const importNameInfo = [ ] ; for ( const decl of result . declarations . keys ( ) ?? [ ] ) { const importInfo = this . _getImportInfo ( decl , filePath ) ; if ( ! importInfo ) { continue ; } const moduleAndType = this . _importResolver . getModuleNameForImport ( importInfo . filePath , execEnv ) ; if ( ! moduleAndType . moduleName ) { if ( ! importInfo . nameInfo . name ) { continue ; } moduleAndType . moduleName = getRelativeModuleName ( this . _importResolver . fileSystem , filePath , importInfo . filePath ) ; } addIfUnique ( importNameInfo , { module : moduleAndType , name : importInfo . nameInfo . name , alias : importInfo . nameInfo . alias } , ( a , b ) => this . _areSame ( a , b ) ) ; } const edits = [ ] ; const newNameInfo = [ ] ; for ( const moduleAndInfo of createMapFromItems ( importNameInfo , ( i ) => i . module . moduleName ) ) { if ( ! this . _tryProcessExistingImports ( moduleAndInfo , importStatements , parseResults , edits ) ) { newNameInfo . push ( ... moduleAndInfo [ 1 ] ) ; continue ; } } edits . push ( ... getTextEditsForAutoImportInsertions ( newNameInfo , importStatements , parseResults , convertOffsetToPosition ( insertionPosition , parseResults . tokenizerOutput . lines ) ) ) ; return edits ; } private _tryProcessExistingImports ( moduleAndInfo , importStatements , parseResults , edits ) { for ( const kindAndImports of createMapFromItems ( importStatements . orderedImports . filter ( ( i ) => i . moduleName === moduleAndInfo [ 0 ] ) , ( i ) => ( i . node . nodeType === ParseNodeType . Import ? 'import' : 'importFrom' ) ) ) { if ( kindAndImports [ 0 ] === 'importFrom' ) { const imported = kindAndImports [ 1 ] . filter ( ( i ) => i . node . nodeType === ParseNodeType . ImportFrom && ! i . node . isWildcardImport ) ; if ( imported . length === 0 ) { continue ; } const info = moduleAndInfo [ 1 ] . filter ( ( m ) => ! imported . some ( ( n ) => n . node . nodeType === ParseNodeType . ImportFrom && n . node . imports . some ( ( i ) => i . name . value === m . name && i . alias ?. value === m . alias ) ) ) ; edits . push ( ... getTextEditsForAutoImportSymbolAddition ( info , imported [ 0 ] , parseResults ) ) ; return true ; } if ( kindAndImports [ 0 ] === 'import' ) { removeArrayElements ( moduleAndInfo [ 1 ] , ( i ) => ! i . name && kindAndImports [ 1 ] . some ( ( n ) => i . alias === n . subnode ?. alias ?. value ) ) ; continue ; } } return false ; } private _getImportInfo ( decl , destFilePath ) { if ( isAliasDeclaration ( decl ) ) { if ( ! decl . node ) { return { filePath : decl . path , nameInfo : { } } ; } if ( decl . node . nodeType === ParseNodeType . ImportAs ) { const importDecl = this . _evaluator . getDeclarationsForNameNode ( decl . node . module . nameParts [ decl . node . module . nameParts . length - 1 ] ) ; if ( ! importDecl || importDecl . length === 0 ) { return undefined ; } return { filePath : importDecl [ 0 ] . path , nameInfo : { alias : decl . usesLocalName ? decl . node . alias ?. value : undefined } , } ; } if ( decl . node . nodeType === ParseNodeType . ImportFromAs ) { let path = decl . path ; if ( ! path ) { if ( decl . submoduleFallback ) { path = getDirectoryPath ( decl . submoduleFallback . path ) ; } } if ( ! path ) { return undefined ; } if ( path === destFilePath && ! decl . usesLocalName && ! decl . submoduleFallback ) { return undefined ; } return { filePath : path , nameInfo : { name : decl . symbolName , alias : decl . usesLocalName ? decl . node . alias ?. value : undefined , } , } ; } if ( decl . node . nodeType === ParseNodeType . ImportFrom ) { return { filePath : decl . path , nameInfo : { name : decl . symbolName } , } ; } } if ( isVariableDeclaration ( decl ) || isFunctionDeclaration ( decl ) || isClassDeclaration ( decl ) ) { const name = getNameFromDeclaration ( decl ) ; if ( ! name ) { return undefined ; } return { filePath : decl . path , nameInfo : { name } , } ; } return undefined ; } private _areSame ( a , b ) { return ( a . alias === b . alias && a . name === b . name && a . module . importType === b . module . importType && a . module . isLocalTypingsFile === b . module . isLocalTypingsFile && a . module . moduleName === b . module . moduleName ) ; } } class NameCollector extends ParseTreeWalker { private readonly _filePath : string ; readonly declsForSymbols = new Map < Declaration , NameNode [ ] > ( ) ; containsUnreferenceableSymbols = false ; constructor ( private _evaluator , private _parseResults , private _range , private _token ) { super ( ) ; this . _filePath = getFileInfo ( this . _parseResults . parseTree ) . filePath ; } override walk ( node ) { if ( ! TextRange . overlapsRange ( this . _range , node ) ) { return ; } super . walk ( node ) ; } override visitModuleName ( node ) { return false ; } override visitName ( name ) { throwIfCancellationRequested ( this . _token ) ; if ( ! isLastNameOfDottedName ( name ) ) { return false ; } const dottedName = getDottedName ( getDottedNameWithGivenNodeAsLastName ( name ) ) ; if ( ! dottedName ) { return false ; } const firstName = dottedName [ 0 ] ; const firstNameDecls = this . _getDeclarationsInModule ( firstName ) ; if ( ! firstNameDecls || firstNameDecls . length === 0 ) { return false ; } if ( dottedName . length === 1 ) { this . _handleName ( firstName , firstNameDecls ) ; return false ; } for ( const firstNameDecl of firstNameDecls ) { if ( ! isAliasDeclaration ( firstNameDecl ) || firstNameDecl . node . nodeType !== ParseNodeType . ImportAs ) { this . _handleName ( firstName , [ firstNameDecl ] ) ; continue ; } if ( firstNameDecl . usesLocalName ) { this . _handleName ( firstName , [ firstNameDecl ] ) ; continue ; } this . _handleImplicitImports ( firstNameDecl , dottedName , 1 ) ; } return false ; } private _getDeclarationsInModule ( name ) { return this . _evaluator . getDeclarationsForNameNode ( name ) ?. filter ( ( d ) => isDefinedInFile ( d , this . _filePath ) ) ; } private _handleImplicitImports ( aliasDecl , dottedName , nameIndex ) { if ( dottedName . length === nameIndex ) { return ; } if ( ! aliasDecl . implicitImports ) { this . _handleName ( dottedName [ nameIndex - 1 ] , [ createSynthesizedAliasDeclaration ( aliasDecl . path ) ] ) ; return ; } const implicitImportDecl = aliasDecl . implicitImports . get ( dottedName [ nameIndex ] . value ) ; if ( ! implicitImportDecl ) { this . _handleName ( dottedName [ nameIndex - 1 ] , [ createSynthesizedAliasDeclaration ( aliasDecl . path ) ] ) ; return ; } this . _handleImplicitImports ( implicitImportDecl , dottedName , nameIndex + 1 ) ; } private _handleName ( name , decls ) { for ( const decl of decls ) { if ( decl . node && TextRange . containsRange ( this . _range , decl . node ) ) { continue ; } if ( isParameterDeclaration ( decl ) ) { this . containsUnreferenceableSymbols = true ; continue ; } if ( isVariableDeclaration ( decl ) || isFunctionDeclaration ( decl ) || isClassDeclaration ( decl ) ) { const scope = getScopeForNode ( name ) ; if ( ! scope ) { this . containsUnreferenceableSymbols = true ; continue ; } const result = scope . lookUpSymbolRecursive ( name . value ) ; if ( ! result || result . scope . type !== ScopeType . Module ) { this . containsUnreferenceableSymbols = true ; continue ; } } this . _addName ( decl , name ) ; } } private _addName ( decl , name ) { getOrAdd ( this . declsForSymbols , decl , ( ) => [ ] ) . push ( name ) ; } }