import { CancellationToken } from 'vscode-languageserver' ; import * as AnalyzerNodeInfo from '../analyzer/analyzerNodeInfo' ; import { AliasDeclaration , Declaration , DeclarationType , isAliasDeclaration } from '../analyzer/declaration' ; import { areDeclarationsSame , createSynthesizedAliasDeclaration , getDeclarationsWithUsesLocalNameRemoved , } from '../analyzer/declarationUtils' ; import { getModuleNode , getStringNodeValueRange } from '../analyzer/parseTreeUtils' ; import { ParseTreeWalker } from '../analyzer/parseTreeWalker' ; import * as ScopeUtils from '../analyzer/scopeUtils' ; import { isStubFile , SourceMapper } from '../analyzer/sourceMapper' ; import { TypeEvaluator } from '../analyzer/typeEvaluatorTypes' ; import { TypeCategory } from '../analyzer/types' ; import { throwIfCancellationRequested } from '../common/cancellationUtils' ; import { TextRange } from '../common/textRange' ; import { ImportAsNode , NameNode , ParseNode , ParseNodeType , StringNode } from '../parser/parseNodes' ; export type CollectionResult = { node : NameNode | StringNode ; range : TextRange ; } ; export class DocumentSymbolCollector extends ParseTreeWalker { static collectFromNode ( node , evaluator , cancellationToken , startingNode ? , treatModuleInImportAndFromImportSame = false ) { const symbolName = node . value ; const declarations = this . getDeclarationsForNode ( node , evaluator , true , cancellationToken ) ; startingNode = startingNode ?? getModuleNode ( node ) ; if ( ! startingNode ) { return [ ] ; } const collector = new DocumentSymbolCollector ( symbolName , declarations , evaluator , cancellationToken , startingNode , treatModuleInImportAndFromImportSame ) ; return collector . collect ( ) ; } static getDeclarationsForNode ( node , evaluator , resolveLocalName , token , sourceMapper ? ) { throwIfCancellationRequested ( token ) ; const declarations = this . _getDeclarationsForNode ( node , evaluator ) ; const resolvedDeclarations = [ ] ; declarations . forEach ( ( decl ) => { const resolvedDecl = evaluator . resolveAliasDeclaration ( decl , resolveLocalName ) ; if ( resolvedDecl ) { resolvedDeclarations . push ( resolvedDecl ) ; if ( sourceMapper && isStubFile ( resolvedDecl . path ) ) { const implDecls = sourceMapper . findDeclarations ( resolvedDecl ) ; for ( const implDecl of implDecls ) { if ( implDecl && implDecl . path ) { this . _addIfUnique ( resolvedDeclarations , implDecl ) ; } } } } } ) ; return resolvedDeclarations ; } private _results : CollectionResult [ ] = [ ] ; private _dunderAllNameNodes = new Set < StringNode > ( ) ; constructor ( private _symbolName , private _declarations , private _evaluator , private _cancellationToken , private _startingNode , private _treatModuleInImportAndFromImportSame = false ) { super ( ) ; this . _setDunderAllNodes ( this . _startingNode ) ; } collect ( ) { this . walk ( this . _startingNode ) ; return this . _results ; } override walk ( node ) { if ( ! AnalyzerNodeInfo . isCodeUnreachable ( node ) ) { super . walk ( node ) ; } } override visitName ( node ) { throwIfCancellationRequested ( this . _cancellationToken ) ; if ( node . value !== this . _symbolName ) { return false ; } if ( this . _declarations . length > 0 ) { const declarations = DocumentSymbolCollector . _getDeclarationsForNode ( node , this . _evaluator ) ; if ( declarations && declarations . length > 0 ) { if ( declarations . some ( ( decl ) => this . _resultsContainsDeclaration ( decl ) ) ) { this . _addResult ( node ) ; } } } else { this . _addResult ( node ) ; } return false ; } override visitString ( node ) { throwIfCancellationRequested ( this . _cancellationToken ) ; if ( this . _dunderAllNameNodes . has ( node ) ) { this . _addResult ( node ) ; } return false ; } private _addResult ( node ) { const range = node . nodeType === ParseNodeType . Name ? node : getStringNodeValueRange ( node ) ; this . _results . push ( { node , range } ) ; } private _resultsContainsDeclaration ( declaration ) { const resolvedDecl = this . _evaluator . resolveAliasDeclaration ( declaration , false ) ; if ( ! resolvedDecl ) { return false ; } if ( this . _declarations . some ( ( decl ) => areDeclarationsSame ( decl , resolvedDecl , this . _treatModuleInImportAndFromImportSame ) ) ) { return true ; } const resolvedDeclNonlocal = this . _getResolveAliasDeclaration ( resolvedDecl ) ; if ( ! resolvedDeclNonlocal || resolvedDeclNonlocal === resolvedDecl ) { return false ; } return this . _declarations . some ( ( decl ) => areDeclarationsSame ( decl , resolvedDeclNonlocal , this . _treatModuleInImportAndFromImportSame ) ) ; } private _getResolveAliasDeclaration ( declaration ) { if ( isAliasDeclFromImportAsWithAlias ( declaration ) ) { return getDeclarationsWithUsesLocalNameRemoved ( [ declaration ] ) [ 0 ] ; } const resolvedDecl = this . _evaluator . resolveAliasDeclaration ( declaration , true ) ; return isAliasDeclFromImportAsWithAlias ( resolvedDecl ) ? getDeclarationsWithUsesLocalNameRemoved ( [ resolvedDecl ] ) [ 0 ] : resolvedDecl ; function isAliasDeclFromImportAsWithAlias ( decl ? ) { return ( ! ! decl && decl . type === DeclarationType . Alias && decl . node && decl . usesLocalName && decl . node . nodeType === ParseNodeType . ImportAs ) ; } } private _setDunderAllNodes ( node ) { if ( node . nodeType !== ParseNodeType . Module ) { return ; } const dunderAllInfo = AnalyzerNodeInfo . getDunderAllInfo ( node ) ; if ( ! dunderAllInfo ) { return ; } const moduleScope = ScopeUtils . getScopeForNode ( node ) ; if ( ! moduleScope ) { return ; } dunderAllInfo . stringNodes . forEach ( ( stringNode ) => { if ( stringNode . value !== this . _symbolName ) { return ; } const symbolInScope = moduleScope . lookUpSymbolRecursive ( stringNode . value ) ; if ( ! symbolInScope ) { return ; } if ( ! symbolInScope . symbol . getDeclarations ( ) . some ( ( d ) => this . _resultsContainsDeclaration ( d ) ) ) { return ; } this . _dunderAllNameNodes . add ( stringNode ) ; } ) ; } private static _addIfUnique ( declarations , itemToAdd ) { for ( const def of declarations ) { if ( areDeclarationsSame ( def , itemToAdd ) ) { return ; } } declarations . push ( itemToAdd ) ; } private static _getDeclarationsForNode ( node , evaluator ) { if ( node . parent ?. nodeType !== ParseNodeType . ModuleName ) { let decls = evaluator . getDeclarationsForNameNode ( node ) || [ ] ; if ( node . parent ?. nodeType === ParseNodeType . ImportFromAs ) { decls = decls . filter ( ( d ) => d . node === node . parent ) ; } if ( decls . length === 0 ) { const type = evaluator . getType ( node ) ; if ( type ?. category === TypeCategory . Module ) { return [ createSynthesizedAliasDeclaration ( type . filePath ) ] ; } } for ( const aliasDecl of decls . filter ( ( d ) => isAliasDeclaration ( d ) && ! d . loadSymbolsFromPath ) ) { const node = ( aliasDecl as AliasDeclaration ) . node ; if ( node . nodeType === ParseNodeType . ImportFromAs ) { continue ; } decls . push ( ... ( evaluator . getDeclarationsForNameNode ( node . module . nameParts [ 0 ] ) || [ ] ) ) ; } return decls ; } const moduleName = node . parent ; if ( moduleName . parent ?. nodeType === ParseNodeType . ImportAs || moduleName . parent ?. nodeType === ParseNodeType . ImportFrom ) { const index = moduleName . nameParts . findIndex ( ( n ) => n === node ) ; if ( index === 0 ) { let decls = [ ] ; const isImportAsWithAlias = moduleName . nameParts . length === 1 && moduleName . parent . nodeType === ParseNodeType . ImportAs && ! ! moduleName . parent . alias ; const importName = isImportAsWithAlias ? ( moduleName . parent as ImportAsNode ) . alias ! . value : moduleName . nameParts [ 0 ] . value ; const symbolWithScope = ScopeUtils . getScopeForNode ( node ) ?. lookUpSymbolRecursive ( importName ) ; if ( symbolWithScope && moduleName . nameParts . length === 1 ) { decls . push ( ... symbolWithScope . symbol . getDeclarations ( ) . filter ( ( d ) => isAliasDeclaration ( d ) ) ) ; if ( decls . length > 1 ) { decls = decls . filter ( ( d ) => { d = d as AliasDeclaration ; if ( d . firstNamePart !== undefined ) { return d . firstNamePart === moduleName . nameParts [ 0 ] . value ; } return d . node === moduleName . parent ; } ) ; } if ( isImportAsWithAlias ) { decls = getDeclarationsWithUsesLocalNameRemoved ( decls ) ; } } decls . push ( ... ( evaluator . getDeclarationsForNameNode ( moduleName . nameParts [ 0 ] ) ?. filter ( ( d ) => isAliasDeclaration ( d ) ) || [ ] ) ) ; return decls ; } if ( index > 0 ) { return evaluator . getDeclarationsForNameNode ( moduleName . nameParts [ index ] ) || [ ] ; } return [ ] ; } return [ ] ; } }