import * as AnalyzerNodeInfo from '../analyzer/analyzerNodeInfo' ; import * as ParseTreeUtils from '../analyzer/parseTreeUtils' ; import { ExecutionEnvironment } from '../common/configOptions' ; import { isDefined } from '../common/core' ; import { getAnyExtensionFromPath } from '../common/pathUtils' ; import { ClassNode , ModuleNode , ParseNode , ParseNodeType } from '../parser/parseNodes' ; import { ClassDeclaration , Declaration , FunctionDeclaration , isAliasDeclaration , isClassDeclaration , isFunctionDeclaration , isParameterDeclaration , isSpecialBuiltInClassDeclaration , isVariableDeclaration , ParameterDeclaration , SpecialBuiltInClassDeclaration , VariableDeclaration , } from './declaration' ; import { ImportResolver } from './importResolver' ; import { SourceFile } from './sourceFile' ; import { TypeEvaluator } from './typeEvaluatorTypes' ; import { ClassType , isFunction , isInstantiableClass , isOverloadedFunction } from './types' ; import { lookUpClassMember } from './typeUtils' ; type ClassOrFunctionOrVariableDeclaration = ClassDeclaration | FunctionDeclaration | VariableDeclaration ; export type ShadowFileBinder = ( stubFilePath : string , implFilePath : string ) => SourceFile | undefined ; export type BoundSourceGetter = ( filePath : string ) => SourceFile | undefined ; export class SourceMapper { constructor ( private _importResolver , private _execEnv , private _evaluator , private _fileBinder , private _boundSourceGetter , private _mapCompiled , private _preferStubs ) { } findModules ( stubFilePath ) { const sourceFiles = this . _getBoundSourceFilesFromStubFile ( stubFilePath ) ; return sourceFiles . map ( ( sf ) => sf . getParseResults ( ) ?. parseTree ) . filter ( isDefined ) ; } findDeclarations ( stubDecl ) { if ( isClassDeclaration ( stubDecl ) ) { return this . _findClassOrTypeAliasDeclarations ( stubDecl ) ; } else if ( isFunctionDeclaration ( stubDecl ) ) { return this . _findFunctionOrTypeAliasDeclarations ( stubDecl ) ; } else if ( isVariableDeclaration ( stubDecl ) ) { return this . _findVariableDeclarations ( stubDecl ) ; } else if ( isParameterDeclaration ( stubDecl ) ) { return this . _findParameterDeclarations ( stubDecl ) ; } else if ( isSpecialBuiltInClassDeclaration ( stubDecl ) ) { return this . _findSpecialBuiltInClassDeclarations ( stubDecl ) ; } return [ ] ; } findClassDeclarations ( stubDecl ) { return this . _findClassOrTypeAliasDeclarations ( stubDecl ) . filter ( ( d ) => isClassDeclaration ( d ) ) . map ( ( d ) => d as ClassDeclaration ) ; } findClassDeclarationsByType ( originatedPath , type ) { const result = [ ] ; this . _addClassTypeDeclarations ( originatedPath , type , result , new Set < string > ( ) ) ; return result . filter ( ( r ) => isClassDeclaration ( r ) ) . map ( ( r ) => r as ClassDeclaration ) ; } findFunctionDeclarations ( stubDecl ) { return this . _findFunctionOrTypeAliasDeclarations ( stubDecl ) . filter ( ( d ) => isFunctionDeclaration ( d ) ) . map ( ( d ) => d as FunctionDeclaration ) ; } private _findSpecialBuiltInClassDeclarations ( stubDecl , recursiveDeclCache = new Set < string > ( ) ) { if ( stubDecl . node . valueExpression . nodeType === ParseNodeType . Name ) { const className = stubDecl . node . valueExpression . value ; const sourceFiles = this . _getBoundSourceFilesFromStubFile ( stubDecl . path ) ; return sourceFiles . flatMap ( ( sourceFile ) => this . _findClassDeclarationsByName ( sourceFile , className , recursiveDeclCache ) ) ; } return [ ] ; } private _findClassOrTypeAliasDeclarations ( stubDecl , recursiveDeclCache = new Set < string > ( ) ) { const className = this . _getFullClassName ( stubDecl . node ) ; const sourceFiles = this . _getBoundSourceFilesFromStubFile ( stubDecl . path ) ; return sourceFiles . flatMap ( ( sourceFile ) => this . _findClassDeclarationsByName ( sourceFile , className , recursiveDeclCache ) ) ; } private _findFunctionOrTypeAliasDeclarations ( stubDecl , recursiveDeclCache = new Set < string > ( ) ) { const functionName = stubDecl . node . name . value ; const sourceFiles = this . _getBoundSourceFilesFromStubFile ( stubDecl . path ) ; if ( stubDecl . isMethod ) { const classNode = ParseTreeUtils . getEnclosingClass ( stubDecl . node ) ; if ( classNode === undefined ) { return [ ] ; } const className = this . _getFullClassName ( classNode ) ; return sourceFiles . flatMap ( ( sourceFile ) => this . _findMethodDeclarationsByName ( sourceFile , className , functionName , recursiveDeclCache ) ) ; } else { return sourceFiles . flatMap ( ( sourceFile ) => this . _findFunctionDeclarationsByName ( sourceFile , functionName , recursiveDeclCache ) ) ; } } private _findVariableDeclarations ( stubDecl , recursiveDeclCache = new Set < string > ( ) ) { if ( stubDecl . node . nodeType !== ParseNodeType . Name ) { return [ ] ; } const variableName = stubDecl . node . value ; const sourceFiles = this . _getBoundSourceFilesFromStubFile ( stubDecl . path ) ; const classNode = ParseTreeUtils . getEnclosingClass ( stubDecl . node ) ; if ( classNode ) { const className = this . _getFullClassName ( classNode ) ; return sourceFiles . flatMap ( ( sourceFile ) => this . _findFieldDeclarationsByName ( sourceFile , className , variableName , recursiveDeclCache ) ) ; } else { return sourceFiles . flatMap ( ( sourceFile ) => this . _findVariableDeclarationsByName ( sourceFile , variableName , recursiveDeclCache ) ) ; } } private _findParameterDeclarations ( stubDecl ) { const result = [ ] ; if ( ! stubDecl . node . name ) { return result ; } const functionNode = ParseTreeUtils . getEnclosingFunction ( stubDecl . node ) ; if ( ! functionNode ) { return result ; } const functionStubDecls = this . _evaluator . getDeclarationsForNameNode ( functionNode . name ) ; if ( ! functionStubDecls ) { return result ; } const recursiveDeclCache = new Set < string > ( ) ; for ( const functionStubDecl of functionStubDecls ) { for ( const functionDecl of this . _findFunctionOrTypeAliasDeclarations ( functionStubDecl as FunctionDeclaration , recursiveDeclCache ) ) { result . push ( ... this . _lookUpSymbolDeclarations ( functionDecl . node , stubDecl . node . name . value ) . filter ( ( d ) => isParameterDeclaration ( d ) ) . map ( ( d ) => d as ParameterDeclaration ) ) ; } } return result ; } private _findMemberDeclarationsByName < T extends Declaration > ( sourceFile , className , memberName , declAdder , recursiveDeclCache ) { const result = [ ] ; const classDecls = this . _findClassDeclarationsByName ( sourceFile , className , recursiveDeclCache ) ; for ( const classDecl of classDecls . filter ( ( d ) => isClassDeclaration ( d ) ) . map ( ( d ) => d as ClassDeclaration ) ) { const classResults = this . _evaluator . getTypeOfClass ( classDecl . node ) ; if ( ! classResults ) { continue ; } const member = lookUpClassMember ( classResults . classType , memberName ) ; if ( member ) { for ( const decl of member . symbol . getDeclarations ( ) ) { declAdder ( decl , recursiveDeclCache , result ) ; } } } return result ; } private _findFieldDeclarationsByName ( sourceFile , className , variableName , recursiveDeclCache ) { let result = [ ] ; const uniqueId = `@${ sourceFile . getFilePath ( ) }/c/${ className }/v/${ variableName }` ; if ( recursiveDeclCache . has ( uniqueId ) ) { return result ; } recursiveDeclCache . add ( uniqueId ) ; result = this . _findMemberDeclarationsByName ( sourceFile , className , variableName , ( decl , cache , result ) => { if ( isVariableDeclaration ( decl ) ) { if ( this . _isStubThatShouldBeMappedToImplementation ( decl . path ) ) { for ( const implDecl of this . _findVariableDeclarations ( decl , cache ) ) { if ( isVariableDeclaration ( implDecl ) ) { result . push ( implDecl ) ; } } } else { result . push ( decl ) ; } } } , recursiveDeclCache ) ; recursiveDeclCache . delete ( uniqueId ) ; return result ; } private _findMethodDeclarationsByName ( sourceFile , className , functionName , recursiveDeclCache ) { let result = [ ] ; const uniqueId = `@${ sourceFile . getFilePath ( ) }/c/${ className }/f/${ functionName }` ; if ( recursiveDeclCache . has ( uniqueId ) ) { return result ; } recursiveDeclCache . add ( uniqueId ) ; result = this . _findMemberDeclarationsByName ( sourceFile , className , functionName , ( decl , cache , result ) => { if ( isFunctionDeclaration ( decl ) ) { if ( this . _isStubThatShouldBeMappedToImplementation ( decl . path ) ) { result . push ( ... this . _findFunctionOrTypeAliasDeclarations ( decl , cache ) ) ; } else { result . push ( decl ) ; } } } , recursiveDeclCache ) ; recursiveDeclCache . delete ( uniqueId ) ; return result ; } private _findVariableDeclarationsByName ( sourceFile , variableName , recursiveDeclCache ) { const result = [ ] ; const uniqueId = `@${ sourceFile . getFilePath ( ) }/v/${ variableName }` ; if ( recursiveDeclCache . has ( uniqueId ) ) { return result ; } recursiveDeclCache . add ( uniqueId ) ; const moduleNode = sourceFile . getParseResults ( ) ?. parseTree ; if ( ! moduleNode ) { return result ; } const decls = this . _lookUpSymbolDeclarations ( moduleNode , variableName ) ; if ( decls . length === 0 ) { this . _addDeclarationsFollowingWildcardImports ( moduleNode , variableName , result , recursiveDeclCache ) ; } else { for ( const decl of decls ) { this . _addVariableDeclarations ( decl , result , recursiveDeclCache ) ; } } recursiveDeclCache . delete ( uniqueId ) ; return result ; } private _findFunctionDeclarationsByName ( sourceFile , functionName , recursiveDeclCache ) { const result = [ ] ; const uniqueId = `@${ sourceFile . getFilePath ( ) }/f/${ functionName }` ; if ( recursiveDeclCache . has ( uniqueId ) ) { return result ; } recursiveDeclCache . add ( uniqueId ) ; const moduleNode = sourceFile . getParseResults ( ) ?. parseTree ; if ( ! moduleNode ) { return result ; } const decls = this . _lookUpSymbolDeclarations ( moduleNode , functionName ) ; if ( decls . length === 0 ) { this . _addDeclarationsFollowingWildcardImports ( moduleNode , functionName , result , recursiveDeclCache ) ; } else { for ( const decl of decls ) { this . _addClassOrFunctionDeclarations ( decl , result , recursiveDeclCache ) ; } } recursiveDeclCache . delete ( uniqueId ) ; return result ; } private _findClassDeclarationsByName ( sourceFile , fullClassName , recursiveDeclCache ) { let classDecls = [ ] ; const parentNode = sourceFile . getParseResults ( ) ?. parseTree ; if ( parentNode ) { let classNameParts = fullClassName . split ( '.' ) ; if ( classNameParts . length > 0 ) { classDecls = this . _findClassDeclarations ( sourceFile , classNameParts [ 0 ] , parentNode , recursiveDeclCache ) ; classNameParts = classNameParts . slice ( 1 ) ; } for ( const classNamePart of classNameParts ) { classDecls = classDecls . flatMap ( ( parentDecl ) => this . _findClassDeclarations ( sourceFile , classNamePart , parentDecl . node , recursiveDeclCache ) ) ; } } return classDecls ; } private _findClassDeclarations ( sourceFile , className , parentNode , recursiveDeclCache ) { const result = [ ] ; const uniqueId = `@${ sourceFile . getFilePath ( ) }[${ parentNode . start }]${ className }` ; if ( recursiveDeclCache . has ( uniqueId ) ) { return result ; } recursiveDeclCache . add ( uniqueId ) ; const decls = this . _lookUpSymbolDeclarations ( parentNode , className ) ; if ( decls . length === 0 && parentNode . nodeType === ParseNodeType . Module ) { this . _addDeclarationsFollowingWildcardImports ( parentNode , className , result , recursiveDeclCache ) ; } else { for ( const decl of decls ) { this . _addClassOrFunctionDeclarations ( decl , result , recursiveDeclCache ) ; } } recursiveDeclCache . delete ( uniqueId ) ; return result ; } private _addVariableDeclarations ( decl , result , recursiveDeclCache ) { if ( isVariableDeclaration ( decl ) ) { if ( this . _isStubThatShouldBeMappedToImplementation ( decl . path ) ) { result . push ( ... this . _findVariableDeclarations ( decl , recursiveDeclCache ) ) ; } else { result . push ( decl ) ; } } else if ( isAliasDeclaration ( decl ) ) { const resolvedDecl = this . _evaluator . resolveAliasDeclaration ( decl , true ) ; if ( resolvedDecl ) { if ( isVariableDeclaration ( resolvedDecl ) ) { this . _addVariableDeclarations ( resolvedDecl , result , recursiveDeclCache ) ; } else if ( isClassDeclaration ( resolvedDecl ) || isFunctionDeclaration ( resolvedDecl ) ) { this . _addClassOrFunctionDeclarations ( resolvedDecl , result , recursiveDeclCache ) ; } } } } private _addClassOrFunctionDeclarations ( decl , result , recursiveDeclCache ) { if ( isClassDeclaration ( decl ) ) { if ( this . _isStubThatShouldBeMappedToImplementation ( decl . path ) ) { result . push ( ... this . _findClassOrTypeAliasDeclarations ( decl , recursiveDeclCache ) ) ; } else { result . push ( decl ) ; } } else if ( isFunctionDeclaration ( decl ) ) { if ( this . _isStubThatShouldBeMappedToImplementation ( decl . path ) ) { result . push ( ... this . _findFunctionOrTypeAliasDeclarations ( decl , recursiveDeclCache ) ) ; } else { result . push ( decl ) ; } } else if ( isAliasDeclaration ( decl ) ) { const resolvedDecl = this . _evaluator . resolveAliasDeclaration ( decl , true ) ; if ( resolvedDecl && ! isAliasDeclaration ( resolvedDecl ) ) { this . _addClassOrFunctionDeclarations ( resolvedDecl , result , recursiveDeclCache ) ; } } else if ( isVariableDeclaration ( decl ) ) { this . _addVariableDeclarations ( decl , result , recursiveDeclCache ) ; const nodeToBind = decl . typeAliasName ?? decl . node ; const type = this . _evaluator . getType ( nodeToBind ) ; if ( ! type ) { return ; } if ( isFunction ( type ) && type . details . declaration ) { this . _addClassOrFunctionDeclarations ( type . details . declaration , result , recursiveDeclCache ) ; } else if ( isOverloadedFunction ( type ) ) { for ( const overloadDecl of type . overloads . map ( ( o ) => o . details . declaration ) . filter ( isDefined ) ) { this . _addClassOrFunctionDeclarations ( overloadDecl , result , recursiveDeclCache ) ; } } else if ( isInstantiableClass ( type ) ) { this . _addClassTypeDeclarations ( decl . path , type , result , recursiveDeclCache ) ; } } } private _addClassTypeDeclarations ( originated , type , result , recursiveDeclCache ) { const importResult = this . _importResolver . resolveImport ( originated , this . _execEnv , { leadingDots : 0 , nameParts : type . details . moduleName . split ( '.' ) , importedSymbols : [ ] , } ) ; if ( importResult . isImportFound && importResult . resolvedPaths . length > 0 ) { const filePath = importResult . resolvedPaths [ importResult . resolvedPaths . length - 1 ] ; const sourceFiles = this . _getSourceFiles ( filePath ) ; const fullClassName = type . details . fullName . substring ( type . details . moduleName . length + 1 ) ; for ( const sourceFile of sourceFiles ) { result . push ( ... this . _findClassDeclarationsByName ( sourceFile , fullClassName , recursiveDeclCache ) ) ; } } } private _getSourceFiles ( filePath ) { const sourceFiles = [ ] ; if ( this . _isStubThatShouldBeMappedToImplementation ( filePath ) ) { sourceFiles . push ( ... this . _getBoundSourceFilesFromStubFile ( filePath ) ) ; } else { const sourceFile = this . _boundSourceGetter ( filePath ) ; if ( sourceFile ) { sourceFiles . push ( sourceFile ) ; } } return sourceFiles ; } private _addDeclarationsFollowingWildcardImports ( moduleNode , symbolName , result , recursiveDeclCache ) { const table = AnalyzerNodeInfo . getScope ( moduleNode ) ?. symbolTable ; if ( ! table ) { return ; } for ( const symbol of table . values ( ) ) { for ( const decl of symbol . getDeclarations ( ) ) { if ( ! isAliasDeclaration ( decl ) || ! decl . path || decl . node . nodeType !== ParseNodeType . ImportFrom || ! decl . node . isWildcardImport ) { continue ; } const uniqueId = `@${ decl . path }/l/${ symbolName }` ; if ( recursiveDeclCache . has ( uniqueId ) ) { continue ; } recursiveDeclCache . add ( uniqueId ) ; const sourceFiles = this . _getSourceFiles ( decl . path ) ; for ( const sourceFile of sourceFiles ) { const moduleNode = sourceFile . getParseResults ( ) ?. parseTree ; if ( ! moduleNode ) { continue ; } const decls = this . _lookUpSymbolDeclarations ( moduleNode , symbolName ) ; if ( decls . length === 0 ) { this . _addDeclarationsFollowingWildcardImports ( moduleNode , symbolName , result , recursiveDeclCache ) ; } else { for ( const decl of decls ) { const resolvedDecl = this . _evaluator . resolveAliasDeclaration ( decl , true ) ; if ( ! resolvedDecl ) { continue ; } if ( isFunctionDeclaration ( resolvedDecl ) || isClassDeclaration ( resolvedDecl ) ) { this . _addClassOrFunctionDeclarations ( resolvedDecl , result , recursiveDeclCache ) ; } else if ( isVariableDeclaration ( resolvedDecl ) ) { this . _addVariableDeclarations ( resolvedDecl , result , recursiveDeclCache ) ; } } } } } } } private _lookUpSymbolDeclarations ( node , symbolName ) { if ( node === undefined ) { return [ ] ; } const containingScope = AnalyzerNodeInfo . getScope ( node ) ; const symbol = containingScope ?. lookUpSymbol ( symbolName ) ; const decls = symbol ?. getDeclarations ( ) ; return decls ?? [ ] ; } private _getFullClassName ( node ) { const fullName = [ ] ; let current = node ; while ( current !== undefined ) { fullName . push ( current . name . value ) ; current = ParseTreeUtils . getEnclosingClass ( current ) ; } return fullName . reverse ( ) . join ( '.' ) ; } private _getBoundSourceFilesFromStubFile ( stubFilePath ) { const paths = this . _importResolver . getSourceFilesFromStub ( stubFilePath , this . _execEnv , this . _mapCompiled ) ; return paths . map ( ( fp ) => this . _fileBinder ( stubFilePath , fp ) ) . filter ( isDefined ) ; } private _isStubThatShouldBeMappedToImplementation ( filePath ) { if ( this . _preferStubs ) { return false ; } const stub = isStubFile ( filePath ) ; if ( ! stub ) { return false ; } return this . _importResolver . getSourceFilesFromStub ( filePath , this . _execEnv , this . _mapCompiled ) . every ( ( f ) => f !== filePath ) ; } } export function isStubFile ( filePath ) { return getAnyExtensionFromPath ( filePath , [ '.pyi' ] , false ) === '.pyi' ; }