Streaming references/symbols, analysis error, indexing changes (#1093)

.github/workflows/validation.yml

@@ -26,11 +26,15 @@ jobs:
         with:
           node-version: ${{ env.NODE_VERSION }}
 
+      # Don't cache on Windows; the cache ends up being very large and
+      # the Windows implementation of the cache task uses a much slower archiver.
       - name: Get npm cache directory
+        if: runner.os != 'Windows'
         id: npm-cache
         run: |
           echo "::set-output name=dir::$(npm config get cache)"
       - uses: actions/cache@v2
+        if: runner.os != 'Windows'
         with:
           path: ${{ steps.npm-cache.outputs.dir }}
           key: ${{ runner.os }}-node-${{ hashFiles('**/package-lock.json') }}

packages/pyright-internal/src/analyzer/analysis.ts

@@ -28,6 +28,7 @@ export interface AnalysisResults {
     fatalErrorOccurred: boolean;
     configParseErrorOccurred: boolean;
     elapsedTime: number;
+    error?: Error;
 }
 
 export type AnalysisCompleteCallback = (results: AnalysisResults) => void;
@@ -85,6 +86,7 @@ export function analyzeProgram(
             fatalErrorOccurred: true,
             configParseErrorOccurred: false,
             elapsedTime: 0,
+            error: debug.getSerializableError(e),
         });
     }
 

packages/pyright-internal/src/analyzer/importResolver.ts

@@ -28,7 +28,7 @@ import {
     stripFileExtension,
     stripTrailingDirectorySeparator,
 } from '../common/pathUtils';
-import { versionToString } from '../common/pythonVersion';
+import { getPythonVersionStrings } from '../common/pythonVersion';
 import { equateStringsCaseInsensitive } from '../common/stringUtils';
 import * as StringUtils from '../common/stringUtils';
 import { isIdentifierChar, isIdentifierStartChar } from '../parser/characters';
@@ -396,17 +396,16 @@ export class ImportResolver {
         const importFailureInfo: string[] = [];
         const roots = [];
 
-        const pythonVersion = execEnv.pythonVersion;
+        const versionFolders = getPythonVersionStrings(execEnv.pythonVersion);
-        const minorVersion = pythonVersion & 0xff;
-        const versionFolders = ['2and3', '3'];
-        if (minorVersion > 0) {
-            versionFolders.push(versionToString(0x300 + minorVersion));
-        }
-
         const stdTypeshed = this._getTypeshedPath(true, execEnv, importFailureInfo);
         if (stdTypeshed) {
             if (useTypeshedVersionedFolders) {
-                roots.push(...versionFolders.map((vf) => combinePaths(stdTypeshed, vf)));
+                for (const version of versionFolders) {
+                    const path = combinePaths(stdTypeshed, version);
+                    if (this.fileSystem.existsSync(path)) {
+                        roots.push(path);
+                    }
+                }
             } else {
                 roots.push(stdTypeshed);
             }
@@ -422,7 +421,12 @@ export class ImportResolver {
         const typeshedPath = this._getTypeshedPath(false, execEnv, importFailureInfo);
         if (typeshedPath) {
             if (useTypeshedVersionedFolders) {
-                roots.push(...versionFolders.map((vf) => combinePaths(typeshedPath, vf)));
+                for (const version of versionFolders) {
+                    const path = combinePaths(typeshedPath, version);
+                    if (this.fileSystem.existsSync(path)) {
+                        roots.push(path);
+                    }
+                }
             } else {
                 roots.push(typeshedPath);
             }
@@ -976,13 +980,7 @@ export class ImportResolver {
             return undefined;
         }
 
-        const pythonVersion = execEnv.pythonVersion;
+        for (const pythonVersionString of getPythonVersionStrings(execEnv.pythonVersion)) {
-        let minorVersion = pythonVersion & 0xff;
-
-        // Search for module starting at "3.x" down to "3.1", then "3", then "2and3".
-        while (true) {
-            const pythonVersionString =
-                minorVersion > 0 ? versionToString(0x300 + minorVersion) : minorVersion === 0 ? '3' : '2and3';
             const testPath = combinePaths(typeshedPath, pythonVersionString);
             if (this.fileSystem.existsSync(testPath)) {
                 const importInfo = this.resolveAbsoluteImport(
@@ -996,12 +994,6 @@ export class ImportResolver {
                     return importInfo;
                 }
             }
-
-            // We use -1 to indicate "2and3", which is searched after "3.0".
-            if (minorVersion === -1) {
-                break;
-            }
-            minorVersion--;
         }
 
         importFailureInfo.push(`Typeshed path not found`);
@@ -1021,23 +1013,11 @@ export class ImportResolver {
             return;
         }
 
-        const pythonVersion = execEnv.pythonVersion;
+        for (const pythonVersionString of getPythonVersionStrings(execEnv.pythonVersion)) {
-        let minorVersion = pythonVersion & 0xff;
-
-        // Search for module starting at "3.x" down to "3.1", then "3", then "2and3".
-        while (true) {
-            const pythonVersionString =
-                minorVersion > 0 ? versionToString(0x300 + minorVersion) : minorVersion === 0 ? '3' : '2and3';
             const testPath = combinePaths(typeshedPath, pythonVersionString);
             if (this.fileSystem.existsSync(testPath)) {
                 this._getCompletionSuggestionsAbsolute(testPath, moduleDescriptor, suggestions, similarityLimit);
             }
-
-            // We use -1 to indicate "2and3", which is searched after "3.0".
-            if (minorVersion === -1) {
-                break;
-            }
-            minorVersion--;
         }
     }
 

packages/pyright-internal/src/analyzer/program.ts

@@ -47,9 +47,9 @@ import {
 } from '../languageService/autoImporter';
 import { CallHierarchyProvider } from '../languageService/callHierarchyProvider';
 import { CompletionResults } from '../languageService/completionProvider';
-import { IndexResults } from '../languageService/documentSymbolProvider';
+import { IndexOptions, IndexResults, WorkspaceSymbolCallback } from '../languageService/documentSymbolProvider';
 import { HoverResults } from '../languageService/hoverProvider';
-import { ReferencesResult } from '../languageService/referencesProvider';
+import { ReferenceCallback, ReferencesResult } from '../languageService/referencesProvider';
 import { SignatureHelpResults } from '../languageService/signatureHelpProvider';
 import { ImportLookupResult } from './analyzerFileInfo';
 import * as AnalyzerNodeInfo from './analyzerNodeInfo';
@@ -142,10 +142,10 @@ export class Program {
         initialConfigOptions: ConfigOptions,
         console?: ConsoleInterface,
         private _extension?: LanguageServiceExtension,
-        logPrefix = 'FG'
+        logTracker?: LogTracker
     ) {
         this._console = console || new StandardConsole();
-        this._logTracker = new LogTracker(console, logPrefix);
+        this._logTracker = logTracker ?? new LogTracker(console, 'FG');
         this._importResolver = initialImportResolver;
         this._configOptions = initialConfigOptions;
         this._createNewEvaluator();
@@ -436,19 +436,24 @@ export class Program {
             return;
         }
 
+        let count = 0;
         return this._runEvaluatorWithCancellationToken(token, () => {
             // Go through all workspace files to create indexing data.
-            // This will cause all files in the workspace to be parsed and bound. We might
+            // This will cause all files in the workspace to be parsed and bound. But
-            // need to drop some of those parse tree and binding info once indexing is done
+            // _handleMemoryHighUsage will make sure we don't OOM
-            // if it didn't exist before.
             for (const sourceFileInfo of this._sourceFileList) {
                 if (!this._isUserCode(sourceFileInfo)) {
                     continue;
                 }
 
                 this._bindFile(sourceFileInfo);
-                const results = sourceFileInfo.sourceFile.index(false, token);
+                const results = sourceFileInfo.sourceFile.index({ indexingForAutoImportMode: false }, token);
                 if (results) {
+                    if (++count > 2000) {
+                        this._console.warn(`Workspace indexing has hit its upper limit: 2000 files`);
+                        return;
+                    }
+
                     callback(sourceFileInfo.sourceFile.getFilePath(), results);
                 }
 
@@ -658,12 +663,12 @@ export class Program {
         return this._evaluator;
     }
 
-    private _parseFile(fileToParse: SourceFileInfo) {
+    private _parseFile(fileToParse: SourceFileInfo, content?: string) {
         if (!this._isFileNeeded(fileToParse) || !fileToParse.sourceFile.isParseRequired()) {
             return;
         }
 
-        if (fileToParse.sourceFile.parse(this._configOptions, this._importResolver)) {
+        if (fileToParse.sourceFile.parse(this._configOptions, this._importResolver, content)) {
             this._parsedFileCount++;
             this._updateSourceFileImports(fileToParse, this._configOptions);
         }
@@ -683,12 +688,12 @@ export class Program {
 
     // Binds the specified file and all of its dependencies, recursively. If
     // it runs out of time, it returns true. If it completes, it returns false.
-    private _bindFile(fileToAnalyze: SourceFileInfo): void {
+    private _bindFile(fileToAnalyze: SourceFileInfo, content?: string): void {
         if (!this._isFileNeeded(fileToAnalyze) || !fileToAnalyze.sourceFile.isBindingRequired()) {
             return;
         }
 
-        this._parseFile(fileToAnalyze);
+        this._parseFile(fileToAnalyze, content);
 
         // We need to parse and bind the builtins import first.
         let builtinsScope: Scope | undefined;
@@ -737,9 +742,16 @@ export class Program {
 
     // Build a map of all modules within this program and the module-
     // level scope that contains the symbol table for the module.
-    private _buildModuleSymbolsMap(sourceFileToExclude: SourceFileInfo, token: CancellationToken): ModuleSymbolMap {
+    private _buildModuleSymbolsMap(
+        sourceFileToExclude: SourceFileInfo,
+        userFileOnly: boolean,
+        token: CancellationToken
+    ): ModuleSymbolMap {
+        // If we have library map, always use the map for library symbols.
         return buildModuleSymbolsMap(
-            this._sourceFileList.filter((s) => s !== sourceFileToExclude),
+            this._sourceFileList.filter(
+                (s) => s !== sourceFileToExclude && (userFileOnly ? this._isUserCode(s) : true)
+            ),
             token
         );
     }
@@ -972,9 +984,9 @@ export class Program {
             }
 
             const writtenWord = fileContents.substr(textRange.start, textRange.length);
-            const map = this._buildModuleSymbolsMap(sourceFileInfo, token);
+            const map = this._buildModuleSymbolsMap(sourceFileInfo, !!libraryMap, token);
             const autoImporter = new AutoImporter(
-                this._configOptions,
+                this._configOptions.findExecEnvironment(filePath),
                 this._importResolver,
                 parseTree,
                 range.start,
@@ -1084,16 +1096,17 @@ export class Program {
         });
     }
 
-    getReferencesForPosition(
+    reportReferencesForPosition(
         filePath: string,
         position: Position,
         includeDeclaration: boolean,
+        reporter: ReferenceCallback,
         token: CancellationToken
-    ): DocumentRange[] | undefined {
+    ) {
-        return this._runEvaluatorWithCancellationToken(token, () => {
+        this._runEvaluatorWithCancellationToken(token, () => {
             const sourceFileInfo = this._getSourceFileInfoFromPath(filePath);
             if (!sourceFileInfo) {
-                return undefined;
+                return;
             }
 
             const invokedFromUserFile = this._isUserCode(sourceFileInfo);
@@ -1104,11 +1117,12 @@ export class Program {
                 this._createSourceMapper(execEnv),
                 position,
                 this._evaluator!,
+                reporter,
                 token
             );
 
             if (!referencesResult) {
-                return undefined;
+                return;
             }
 
             // Do we need to do a global search as well?
@@ -1155,19 +1169,18 @@ export class Program {
                             continue;
                         }
 
-                        const tempResult: ReferencesResult = {
+                        const tempResult = new ReferencesResult(
-                            requiresGlobalSearch: referencesResult.requiresGlobalSearch,
+                            referencesResult.requiresGlobalSearch,
-                            nodeAtOffset: referencesResult.nodeAtOffset,
+                            referencesResult.nodeAtOffset,
-                            symbolName: referencesResult.symbolName,
+                            referencesResult.symbolName,
-                            declarations: referencesResult.declarations,
+                            referencesResult.declarations
-                            locations: [],
+                        );
-                        };
 
                         declFileInfo.sourceFile.addReferences(tempResult, includeDeclaration, this._evaluator!, token);
                         for (const loc of tempResult.locations) {
                             // Include declarations only. And throw away any references
                             if (loc.path === decl.path && doesRangeContain(decl.range, loc.range)) {
-                                referencesResult.locations.push(loc);
+                                referencesResult.addLocations(loc);
                             }
                         }
                     }
@@ -1175,13 +1188,11 @@ export class Program {
             } else {
                 sourceFileInfo.sourceFile.addReferences(referencesResult, includeDeclaration, this._evaluator!, token);
             }
-
-            return referencesResult.locations;
         });
     }
 
-    getFileIndex(filePath: string, importSymbolsOnly: boolean, token: CancellationToken): IndexResults | undefined {
+    getFileIndex(filePath: string, options: IndexOptions, token: CancellationToken): IndexResults | undefined {
-        if (importSymbolsOnly) {
+        if (options.indexingForAutoImportMode) {
             // Memory optimization. We only want to hold onto symbols
             // usable outside when importSymbolsOnly is on.
             const name = stripFileExtension(getFileName(filePath));
@@ -1198,8 +1209,26 @@ export class Program {
                 return undefined;
             }
 
-            this._bindFile(sourceFileInfo);
+            let content: string | undefined = undefined;
-            return sourceFileInfo.sourceFile.index(importSymbolsOnly, token);
+            if (
+                options.indexingForAutoImportMode &&
+                !sourceFileInfo.sourceFile.isStubFile() &&
+                sourceFileInfo.sourceFile.getClientVersion() === null
+            ) {
+                try {
+                    // Perf optimization. if py file doesn't contain __all__
+                    // No need to parse and bind.
+                    content = this._fs.readFileSync(filePath, 'utf8');
+                    if (content.indexOf('__all__') < 0) {
+                        return undefined;
+                    }
+                } catch (error) {
+                    content = undefined;
+                }
+            }
+
+            this._bindFile(sourceFileInfo, content);
+            return sourceFileInfo.sourceFile.index(options, token);
         });
     }
 
@@ -1217,8 +1246,8 @@ export class Program {
         });
     }
 
-    addSymbolsForWorkspace(symbolList: SymbolInformation[], query: string, token: CancellationToken) {
+    reportSymbolsForWorkspace(query: string, reporter: WorkspaceSymbolCallback, token: CancellationToken) {
-        return this._runEvaluatorWithCancellationToken(token, () => {
+        this._runEvaluatorWithCancellationToken(token, () => {
             // Don't do a search if the query is empty. We'll return
             // too many results in this case.
             if (!query) {
@@ -1236,7 +1265,10 @@ export class Program {
                     this._bindFile(sourceFileInfo);
                 }
 
-                sourceFileInfo.sourceFile.addSymbolsForDocument(symbolList, query, token);
+                const symbolList = sourceFileInfo.sourceFile.getSymbolsForDocument(query, token);
+                if (symbolList.length > 0) {
+                    reporter(symbolList);
+                }
 
                 // This operation can consume significant memory, so check
                 // for situations where we need to discard the type cache.
@@ -1334,7 +1366,7 @@ export class Program {
                 this._evaluator!,
                 this._createSourceMapper(execEnv, /* mapCompiled */ true),
                 libraryMap,
-                () => this._buildModuleSymbolsMap(sourceFileInfo, token),
+                () => this._buildModuleSymbolsMap(sourceFileInfo, !!libraryMap, token),
                 token
             );
         });
@@ -1384,7 +1416,7 @@ export class Program {
                 this._evaluator!,
                 this._createSourceMapper(execEnv, /* mapCompiled */ true),
                 libraryMap,
-                () => this._buildModuleSymbolsMap(sourceFileInfo, token),
+                () => this._buildModuleSymbolsMap(sourceFileInfo, !!libraryMap, token),
                 completionItem,
                 token
             );
@@ -1410,6 +1442,7 @@ export class Program {
                 this._createSourceMapper(execEnv),
                 position,
                 this._evaluator!,
+                undefined,
                 token
             );
 
@@ -1472,6 +1505,7 @@ export class Program {
             this._createSourceMapper(execEnv),
             position,
             this._evaluator!,
+            undefined,
             token
         );
 
@@ -1508,6 +1542,7 @@ export class Program {
             this._createSourceMapper(execEnv),
             position,
             this._evaluator!,
+            undefined,
             token
         );
 
@@ -1563,6 +1598,7 @@ export class Program {
             this._createSourceMapper(execEnv),
             position,
             this._evaluator!,
+            undefined,
             token
         );
 

packages/pyright-internal/src/analyzer/service.ts

@@ -46,8 +46,9 @@ import {
 import { DocumentRange, Position, Range } from '../common/textRange';
 import { timingStats } from '../common/timing';
 import { CompletionResults } from '../languageService/completionProvider';
-import { IndexResults } from '../languageService/documentSymbolProvider';
+import { IndexResults, WorkspaceSymbolCallback } from '../languageService/documentSymbolProvider';
 import { HoverResults } from '../languageService/hoverProvider';
+import { ReferenceCallback } from '../languageService/referencesProvider';
 import { SignatureHelpResults } from '../languageService/signatureHelpProvider';
 import { AnalysisCompleteCallback } from './analysis';
 import { BackgroundAnalysisProgram, BackgroundAnalysisProgramFactory } from './backgroundAnalysisProgram';
@@ -244,21 +245,22 @@ export class AnalyzerService {
         return this._program.getDefinitionsForPosition(filePath, position, token);
     }
 
-    getReferencesForPosition(
+    reportReferencesForPosition(
         filePath: string,
         position: Position,
         includeDeclaration: boolean,
+        reporter: ReferenceCallback,
         token: CancellationToken
-    ): DocumentRange[] | undefined {
+    ) {
-        return this._program.getReferencesForPosition(filePath, position, includeDeclaration, token);
+        this._program.reportReferencesForPosition(filePath, position, includeDeclaration, reporter, token);
     }
 
     addSymbolsForDocument(filePath: string, symbolList: DocumentSymbol[], token: CancellationToken) {
         this._program.addSymbolsForDocument(filePath, symbolList, token);
     }
 
-    addSymbolsForWorkspace(symbolList: SymbolInformation[], query: string, token: CancellationToken) {
+    reportSymbolsForWorkspace(query: string, reporter: WorkspaceSymbolCallback, token: CancellationToken) {
-        this._program.addSymbolsForWorkspace(symbolList, query, token);
+        this._program.reportSymbolsForWorkspace(query, reporter, token);
     }
 
     getHoverForPosition(filePath: string, position: Position, token: CancellationToken): HoverResults | undefined {

packages/pyright-internal/src/analyzer/sourceFile.ts

@@ -36,10 +36,10 @@ import { CompletionResults } from '../languageService/completionProvider';
 import { CompletionItemData, CompletionProvider } from '../languageService/completionProvider';
 import { DefinitionProvider } from '../languageService/definitionProvider';
 import { DocumentHighlightProvider } from '../languageService/documentHighlightProvider';
-import { DocumentSymbolProvider, IndexResults } from '../languageService/documentSymbolProvider';
+import { DocumentSymbolProvider, IndexOptions, IndexResults } from '../languageService/documentSymbolProvider';
 import { HoverProvider, HoverResults } from '../languageService/hoverProvider';
 import { performQuickAction } from '../languageService/quickActions';
-import { ReferencesProvider, ReferencesResult } from '../languageService/referencesProvider';
+import { ReferenceCallback, ReferencesProvider, ReferencesResult } from '../languageService/referencesProvider';
 import { SignatureHelpProvider, SignatureHelpResults } from '../languageService/signatureHelpProvider';
 import { Localizer } from '../localization/localize';
 import { ModuleNode } from '../parser/parseNodes';
@@ -508,7 +508,7 @@ export class SourceFile {
     // Parse the file and update the state. Callers should wait for completion
     // (or at least cancel) prior to calling again. It returns true if a parse
     // was required and false if the parse information was up to date already.
-    parse(configOptions: ConfigOptions, importResolver: ImportResolver): boolean {
+    parse(configOptions: ConfigOptions, importResolver: ImportResolver, content?: string): boolean {
         return this._logTracker.log(`parsing: ${this._filePath}`, (logState) => {
             // If the file is already parsed, we can skip.
             if (!this.isParseRequired()) {
@@ -520,14 +520,15 @@ export class SourceFile {
             let fileContents = this._fileContents;
             if (this._clientVersion === null) {
                 try {
-                    timingStats.readFileTime.timeOperation(() => {
+                    const elapsedTime = timingStats.readFileTime.timeOperation(() => {
                         // Read the file's contents.
-                        fileContents = this.fileSystem.readFileSync(this._filePath, 'utf8');
+                        fileContents = content ?? this.fileSystem.readFileSync(this._filePath, 'utf8');
 
                         // Remember the length and hash for comparison purposes.
                         this._lastFileContentLength = fileContents.length;
                         this._lastFileContentHash = StringUtils.hashString(fileContents);
                     });
+                    logState.add(`fs read ${elapsedTime}ms`);
                 } catch (error) {
                     diagSink.addError(`Source file could not be read`, getEmptyRange());
                     fileContents = '';
@@ -634,18 +635,22 @@ export class SourceFile {
         });
     }
 
-    index(importSymbolsOnly: boolean, token: CancellationToken): IndexResults | undefined {
+    index(options: IndexOptions, token: CancellationToken): IndexResults | undefined {
+        return this._logTracker.log(`indexing: ${this._filePath}`, (ls) => {
             // If we have no completed analysis job, there's nothing to do.
             if (!this._parseResults || !this.isIndexingRequired()) {
+                ls.suppress();
                 return undefined;
             }
 
             this._indexingNeeded = false;
-        const symbols = DocumentSymbolProvider.indexSymbols(this._parseResults, importSymbolsOnly, token);
+            const symbols = DocumentSymbolProvider.indexSymbols(this._parseResults, options, token);
+            ls.add(`found ${symbols.length}`);
 
             const name = stripFileExtension(getFileName(this._filePath));
             const privateOrProtected = SymbolNameUtils.isPrivateOrProtectedName(name);
             return { privateOrProtected, symbols };
+        });
     }
 
     getDefinitionsForPosition(
@@ -672,6 +677,7 @@ export class SourceFile {
         sourceMapper: SourceMapper,
         position: Position,
         evaluator: TypeEvaluator,
+        reporter: ReferenceCallback | undefined,
         token: CancellationToken
     ): ReferencesResult | undefined {
         // If we have no completed analysis job, there's nothing to do.
@@ -685,6 +691,7 @@ export class SourceFile {
             this._filePath,
             position,
             evaluator,
+            reporter,
             token
         );
     }
@@ -724,18 +731,17 @@ export class SourceFile {
         );
     }
 
-    addSymbolsForDocument(symbolList: SymbolInformation[], query: string, token: CancellationToken) {
+    getSymbolsForDocument(query: string, token: CancellationToken) {
         // If we have no completed analysis job, there's nothing to do.
         if (!this._parseResults && !this._cachedIndexResults) {
-            return;
+            return [];
         }
 
-        DocumentSymbolProvider.addSymbolsForDocument(
+        return DocumentSymbolProvider.getSymbolsForDocument(
             this.getCachedIndexResults(),
             this._parseResults,
             this._filePath,
             query,
-            symbolList,
             token
         );
     }

packages/pyright-internal/src/analyzer/staticExpressions.ts

@@ -8,7 +8,7 @@
  * (parse node trees).
  */
 
-import { ExecutionEnvironment } from '../common/configOptions';
+import { ExecutionEnvironment, PythonPlatform } from '../common/configOptions';
 import { ExpressionNode, NumberNode, ParseNodeType, TupleNode } from '../parser/parseNodes';
 import { KeywordType, OperatorType } from '../parser/tokenizerTypes';
 
@@ -220,11 +220,11 @@ function _isOsNameInfoExpression(node: ExpressionNode): boolean {
 }
 
 function _getExpectedPlatformNameFromPlatform(execEnv: ExecutionEnvironment): string | undefined {
-    if (execEnv.pythonPlatform === 'Darwin') {
+    if (execEnv.pythonPlatform === PythonPlatform.Darwin) {
         return 'darwin';
-    } else if (execEnv.pythonPlatform === 'Windows') {
+    } else if (execEnv.pythonPlatform === PythonPlatform.Windows) {
         return 'win32';
-    } else if (execEnv.pythonPlatform === 'Linux') {
+    } else if (execEnv.pythonPlatform === PythonPlatform.Linux) {
         return 'linux';
     }
 
@@ -232,11 +232,11 @@ function _getExpectedPlatformNameFromPlatform(execEnv: ExecutionEnvironment): st
 }
 
 function _getExpectedOsNameFromPlatform(execEnv: ExecutionEnvironment): string | undefined {
-    if (execEnv.pythonPlatform === 'Darwin') {
+    if (execEnv.pythonPlatform === PythonPlatform.Darwin) {
         return 'posix';
-    } else if (execEnv.pythonPlatform === 'Windows') {
+    } else if (execEnv.pythonPlatform === PythonPlatform.Windows) {
         return 'nt';
-    } else if (execEnv.pythonPlatform === 'Linux') {
+    } else if (execEnv.pythonPlatform === PythonPlatform.Linux) {
         return 'posix';
     }
 

packages/pyright-internal/src/backgroundAnalysisBase.ts

@@ -32,6 +32,7 @@ import { Diagnostic } from './common/diagnostic';
 import { FileDiagnostics } from './common/diagnosticSink';
 import { LanguageServiceExtension } from './common/extensibility';
 import { FileSystem } from './common/fileSystem';
+import { LogTracker } from './common/logTracker';
 import { Range } from './common/textRange';
 import { IndexResults } from './languageService/documentSymbolProvider';
 
@@ -252,12 +253,13 @@ export class BackgroundAnalysisRunnerBase extends BackgroundThreadBase {
 
         this._configOptions = new ConfigOptions(data.rootDirectory);
         this._importResolver = this.createImportResolver(this.fs, this._configOptions);
+        const console = this.getConsole();
         this._program = new Program(
             this._importResolver,
             this._configOptions,
-            this.getConsole(),
+            console,
             this._extension,
-            `BG(${threadId})`
+            new LogTracker(console, `BG(${threadId})`)
         );
     }
 

packages/pyright-internal/src/common/configOptions.ts

@@ -31,6 +31,12 @@ import {
     versionToString,
 } from './pythonVersion';
 
+export enum PythonPlatform {
+    Darwin = 'Darwin',
+    Windows = 'Windows',
+    Linux = 'Linux',
+}
+
 export class ExecutionEnvironment {
     // Default to "." which indicates every file in the project.
     constructor(root: string, defaultPythonVersion?: PythonVersion, defaultPythonPlatform?: string) {
@@ -1224,11 +1230,11 @@ export class ConfigOptions {
         }
 
         if (process.platform === 'darwin') {
-            this.defaultPythonPlatform = 'Darwin';
+            this.defaultPythonPlatform = PythonPlatform.Darwin;
         } else if (process.platform === 'linux') {
-            this.defaultPythonPlatform = 'Linux';
+            this.defaultPythonPlatform = PythonPlatform.Linux;
         } else if (process.platform === 'win32') {
-            this.defaultPythonPlatform = 'Windows';
+            this.defaultPythonPlatform = PythonPlatform.Windows;
         }
 
         if (this.defaultPythonPlatform !== undefined) {

packages/pyright-internal/src/common/debug.ts

@@ -7,7 +7,7 @@
  */
 
 import { stableSort } from './collectionUtils';
-import { AnyFunction, compareValues, hasProperty } from './core';
+import { AnyFunction, compareValues, hasProperty, isString } from './core';
 
 export function assert(
     expression: boolean,
@@ -106,6 +106,25 @@ export function getErrorString(error: any): string {
     );
 }
 
+export function getSerializableError(error: any): Error | undefined {
+    if (!error) {
+        return undefined;
+    }
+
+    const exception = JSON.stringify(error);
+    if (exception.length > 2) {
+        // Given error object is JSON.stringify serializable. Use it as it is
+        // to preserve properties.
+        return error;
+    }
+
+    // Convert error to JSON.stringify serializable Error shape.
+    const name = error.name ? (isString(error.name) ? error.name : 'noname') : 'noname';
+    const message = error.message ? (isString(error.message) ? error.message : 'nomessage') : 'nomessage';
+    const stack = error.stack ? (isString(error.stack) ? error.stack : undefined) : undefined;
+    return { name, message, stack };
+}
+
 function getEnumMembers(enumObject: any) {
     const result: [number, string][] = [];
     for (const name of Object.keys(enumObject)) {

packages/pyright-internal/src/common/pathUtils.ts

@@ -914,3 +914,14 @@ export function isFileSystemCaseSensitiveInternal(fs: FileSystem) {
         }
     }
 }
+
+export function getLibraryPathWithoutExtension(libraryFilePath: string) {
+    let filePathWithoutExtension = stripFileExtension(libraryFilePath);
+
+    // Strip off the '/__init__' if it's present.
+    if (filePathWithoutExtension.endsWith('__init__')) {
+        filePathWithoutExtension = filePathWithoutExtension.substr(0, filePathWithoutExtension.length - 9);
+    }
+
+    return filePathWithoutExtension;
+}

packages/pyright-internal/src/common/pythonVersion.ts

@@ -71,3 +71,23 @@ export function versionFromMajorMinor(major: number, minor: number): PythonVersi
 export function is3x(version: PythonVersion): boolean {
     return version >> 8 === 3;
 }
+
+export function getPythonVersionStrings(pythonVersion: PythonVersion) {
+    const versionStrings: string[] = [];
+
+    let minorVersion = pythonVersion & 0xff;
+    while (true) {
+        const pythonVersionString =
+            minorVersion > 0 ? versionToString(0x300 + minorVersion) : minorVersion === 0 ? '3' : '2and3';
+        versionStrings.push(pythonVersionString);
+
+        // We use -1 to indicate "2and3", which is searched after "3.0".
+        if (minorVersion === -1) {
+            break;
+        }
+
+        minorVersion--;
+    }
+
+    return versionStrings;
+}

packages/pyright-internal/src/common/timing.ts

@@ -39,8 +39,11 @@ export class TimingStat {
             this.isTiming = true;
             const duration = new Duration();
             callback();
-            this.totalTime += duration.getDurationInMilliseconds();
+            const elapsedTime = duration.getDurationInMilliseconds();
+            this.totalTime += elapsedTime;
             this.isTiming = false;
+
+            return elapsedTime;
         }
     }
 

packages/pyright-internal/src/languageServerBase.ts

@@ -76,10 +76,12 @@ import {
 import { containsPath, convertPathToUri, convertUriToPath } from './common/pathUtils';
 import { ProgressReporter, ProgressReportTracker } from './common/progressReporter';
 import { convertWorkspaceEdits } from './common/textEditUtils';
-import { Position } from './common/textRange';
+import { DocumentRange, Position } from './common/textRange';
 import { AnalyzerServiceExecutor } from './languageService/analyzerServiceExecutor';
 import { CompletionItemData, CompletionResults } from './languageService/completionProvider';
+import { WorkspaceSymbolCallback } from './languageService/documentSymbolProvider';
 import { convertHoverResults } from './languageService/hoverProvider';
+import { ReferenceCallback } from './languageService/referencesProvider';
 import { Localizer } from './localization/localize';
 import { WorkspaceMap } from './workspaceMap';
 
@@ -438,7 +440,7 @@ export abstract class LanguageServerBase implements LanguageServerInterface {
             return locations.map((loc) => Location.create(convertPathToUri(loc.path), loc.range));
         });
 
-        this._connection.onReferences(async (params, token, reporter) => {
+        this._connection.onReferences(async (params, token, workDoneReporter, resultReporter) => {
             if (this._pendingFindAllRefsCancellationSource) {
                 this._pendingFindAllRefsCancellationSource.cancel();
                 this._pendingFindAllRefsCancellationSource = undefined;
@@ -449,9 +451,10 @@ export abstract class LanguageServerBase implements LanguageServerInterface {
             // any long-running actions.
             const progress = await this._getProgressReporter(
                 params.workDoneToken,
-                reporter,
+                workDoneReporter,
                 Localizer.CodeAction.findingReferences()
             );
+
             const source = CancelAfter(token, progress.token);
             this._pendingFindAllRefsCancellationSource = source;
 
@@ -467,18 +470,24 @@ export abstract class LanguageServerBase implements LanguageServerInterface {
                     return;
                 }
 
-                const locations = workspace.serviceInstance.getReferencesForPosition(
+                const convert = (locs: DocumentRange[]): Location[] => {
+                    return locs.map((loc) => Location.create(convertPathToUri(loc.path), loc.range));
+                };
+
+                const locations: Location[] = [];
+                const reporter: ReferenceCallback = resultReporter
+                    ? (locs) => resultReporter.report(convert(locs))
+                    : (locs) => locations.push(...convert(locs));
+
+                workspace.serviceInstance.reportReferencesForPosition(
                     filePath,
                     position,
                     params.context.includeDeclaration,
+                    reporter,
                     source.token
                 );
 
-                if (!locations) {
+                return locations;
-                    return undefined;
-                }
-
-                return locations.map((loc) => Location.create(convertPathToUri(loc.path), loc.range));
             } finally {
                 progress.reporter.done();
                 source.dispose();
@@ -500,13 +509,17 @@ export abstract class LanguageServerBase implements LanguageServerInterface {
             return symbolList;
         });
 
-        this._connection.onWorkspaceSymbol(async (params, token) => {
+        this._connection.onWorkspaceSymbol(async (params, token, _, resultReporter) => {
             const symbolList: SymbolInformation[] = [];
 
+            const reporter: WorkspaceSymbolCallback = resultReporter
+                ? (symbols) => resultReporter.report(symbols)
+                : (symbols) => symbolList.push(...symbols);
+
             for (const workspace of this._workspaceMap.values()) {
                 await workspace.isInitialized.promise;
                 if (!workspace.disableLanguageServices) {
-                    workspace.serviceInstance.addSymbolsForWorkspace(symbolList, params.query, token);
+                    workspace.serviceInstance.reportSymbolsForWorkspace(params.query, reporter, token);
                 }
             }
 

packages/pyright-internal/src/languageService/autoImporter.ts

@@ -7,7 +7,6 @@
 
 import { CancellationToken, CompletionItemKind, SymbolKind } from 'vscode-languageserver';
 
-import * as AnalyzerNodeInfo from '../analyzer/analyzerNodeInfo';
 import { DeclarationType } from '../analyzer/declaration';
 import { ImportResolver, ModuleNameAndType } from '../analyzer/importResolver';
 import { ImportType } from '../analyzer/importResult';
@@ -23,19 +22,14 @@ import { SourceFileInfo } from '../analyzer/program';
 import { Symbol } from '../analyzer/symbol';
 import * as SymbolNameUtils from '../analyzer/symbolNameUtils';
 import { throwIfCancellationRequested } from '../common/cancellationUtils';
-import { ConfigOptions } from '../common/configOptions';
+import { ExecutionEnvironment } from '../common/configOptions';
 import { TextEditAction } from '../common/editAction';
 import { combinePaths, getDirectoryPath, getFileName, stripFileExtension } from '../common/pathUtils';
 import * as StringUtils from '../common/stringUtils';
 import { Position } from '../common/textRange';
 import { ParseNodeType } from '../parser/parseNodes';
 import { ParseResults } from '../parser/parser';
-import {
+import { IndexAliasData, IndexResults } from './documentSymbolProvider';
-    getIndexAliasData,
-    includeAliasDeclarationInIndex,
-    IndexAliasData,
-    IndexResults,
-} from './documentSymbolProvider';
 
 export interface AutoImportSymbol {
     readonly importAlias?: IndexAliasData;
@@ -74,7 +68,6 @@ export function buildModuleSymbolsMap(files: SourceFileInfo[], token: Cancellati
                 return;
             }
 
-            const fileInfo = AnalyzerNodeInfo.getFileInfo(file.sourceFile.getParseResults()!.parseTree);
             moduleSymbolMap.set(filePath, {
                 forEach(callbackfn: (value: AutoImportSymbol, key: string) => void): void {
                     symbolTable.forEach((symbol, name) => {
@@ -92,22 +85,19 @@ export function buildModuleSymbolsMap(files: SourceFileInfo[], token: Cancellati
                             return;
                         }
 
-                        let importAlias: IndexAliasData | undefined;
                         if (declaration.type === DeclarationType.Alias) {
-                            if (!includeAliasDeclarationInIndex(declaration)) {
+                            // We don't include import alias in auto import
+                            // for workspace files.
                             return;
                         }
 
-                            importAlias = getIndexAliasData(fileInfo?.importLookup, declaration);
-                        }
-
                         const variableKind =
                             declaration.type === DeclarationType.Variable &&
                             !declaration.isConstant &&
                             !declaration.isFinal
                                 ? CompletionItemKind.Variable
                                 : undefined;
-                        callbackfn({ importAlias, symbol, kind: variableKind }, name);
+                        callbackfn({ symbol, kind: variableKind }, name);
                     });
                 },
             });
@@ -146,14 +136,14 @@ interface ImportAliasData {
     importParts: ImportParts;
     importGroup: ImportGroup;
     symbol?: Symbol;
+    kind?: CompletionItemKind;
 }
 
 export class AutoImporter {
     private _importStatements: ImportStatements;
-    private _filePath: string;
 
     constructor(
-        private _configOptions: ConfigOptions,
+        private _execEnvironment: ExecutionEnvironment,
         private _importResolver: ImportResolver,
         private _parseResults: ParseResults,
         private _invocationPosition: Position,
@@ -161,7 +151,6 @@ export class AutoImporter {
         private _moduleSymbolMap: ModuleSymbolMap,
         private _libraryMap?: Map<string, IndexResults>
     ) {
-        this._filePath = AnalyzerNodeInfo.getFileInfo(this._parseResults.parseTree)!.filePath;
         this._importStatements = getTopLevelImports(this._parseResults.parseTree);
     }
 
@@ -309,6 +298,7 @@ export class AutoImporter {
                         },
                         importGroup,
                         symbol: autoImportSymbol.symbol,
+                        kind: convertSymbolKindToCompletionItemKind(autoImportSymbol.importAlias.kind),
                     },
                     importAliasMap
                 );
@@ -356,7 +346,7 @@ export class AutoImporter {
         }
 
         this._addToImportAliasMap(
-            { modulePath: filePath, originalName: importParts.importName },
+            { modulePath: filePath, originalName: importParts.importName, kind: SymbolKind.Module },
             { importParts, importGroup },
             importAliasMap
         );
@@ -386,6 +376,7 @@ export class AutoImporter {
                     name: importAliasData.importParts.importName,
                     alias: aliasName,
                     symbol: importAliasData.symbol,
+                    kind: importAliasData.kind,
                     source: importAliasData.importParts.importFrom,
                     edits: autoImportTextEdits,
                 });
@@ -525,8 +516,7 @@ export class AutoImporter {
     // convert to a module name that can be used in an
     // 'import from' statement.
     private _getModuleNameAndTypeFromFilePath(filePath: string): ModuleNameAndType {
-        const execEnvironment = this._configOptions.findExecEnvironment(this._filePath);
+        return this._importResolver.getModuleNameForImport(filePath, this._execEnvironment);
-        return this._importResolver.getModuleNameForImport(filePath, execEnvironment);
     }
 
     private _getImportGroupFromModuleNameAndType(moduleNameAndType: ModuleNameAndType): ImportGroup {

packages/pyright-internal/src/languageService/completionProvider.ts

@@ -965,7 +965,7 @@ export class CompletionProvider {
     private _getAutoImportCompletions(priorWord: string, completionList: CompletionList) {
         const moduleSymbolMap = this._moduleSymbolsCallback();
         const autoImporter = new AutoImporter(
-            this._configOptions,
+            this._configOptions.findExecEnvironment(this._filePath),
             this._importResolver,
             this._parseResults,
             this._position,

packages/pyright-internal/src/languageService/documentSymbolProvider.ts

@@ -20,24 +20,27 @@ import { getLastTypedDeclaredForSymbol } from '../analyzer/symbolUtils';
 import { TypeEvaluator } from '../analyzer/typeEvaluator';
 import { isProperty } from '../analyzer/typeUtils';
 import { throwIfCancellationRequested } from '../common/cancellationUtils';
+import { getLibraryPathWithoutExtension } from '../common/pathUtils';
 import { convertOffsetsToRange } from '../common/positionUtils';
 import * as StringUtils from '../common/stringUtils';
 import { Range } from '../common/textRange';
+import { ModuleNode, ParseNodeType } from '../parser/parseNodes';
 import { ParseResults } from '../parser/parser';
 
 export interface IndexAliasData {
     readonly originalName: string;
     readonly modulePath: string;
+    readonly kind: SymbolKind;
 }
 
 export interface IndexSymbolData {
     readonly name: string;
-    readonly alias: IndexAliasData | undefined;
     readonly externallyVisible: boolean;
     readonly kind: SymbolKind;
-    readonly range: Range;
+    readonly alias?: IndexAliasData;
-    readonly selectionRange: Range;
+    readonly range?: Range;
-    readonly children: IndexSymbolData[];
+    readonly selectionRange?: Range;
+    readonly children?: IndexSymbolData[];
 }
 
 export interface IndexResults {
@@ -45,30 +48,62 @@ export interface IndexResults {
     readonly symbols: IndexSymbolData[];
 }
 
-export function includeAliasDeclarationInIndex(declaration: AliasDeclaration): boolean {
+export interface IndexOptions {
-    return declaration.usesLocalName && !!declaration.symbolName && declaration.path.length > 0;
+    indexingForAutoImportMode: boolean;
+}
+
+export type WorkspaceSymbolCallback = (symbols: SymbolInformation[]) => void;
+
+export function includeAliasDeclarationInIndex(
+    importLookup: ImportLookup,
+    modulePath: string,
+    declaration: AliasDeclaration
+): boolean {
+    const aliasUsed = declaration.usesLocalName && !!declaration.symbolName && declaration.path.length > 0;
+    const wildcardUsed = declaration.node.nodeType === ParseNodeType.ImportFrom && declaration.node.isWildcardImport;
+    if (!aliasUsed && !wildcardUsed) {
+        return false;
+    }
+
+    // Make sure imported symbol is a submodule of same package.
+    if (!getLibraryPathWithoutExtension(declaration.path).startsWith(modulePath)) {
+        return false;
+    }
+
+    if (wildcardUsed) {
+        // if "import *" is used, resolve the alias to see whether we should include it.
+        const resolved = resolveAliasDeclaration(importLookup, declaration, true);
+        if (!resolved) {
+            return false;
+        }
+
+        if (!getLibraryPathWithoutExtension(resolved.path).startsWith(modulePath)) {
+            return false;
+        }
+    }
+
+    return true;
 }
 
 export function getIndexAliasData(
-    importLookup: ImportLookup | undefined,
+    importLookup: ImportLookup,
     declaration: AliasDeclaration
 ): IndexAliasData | undefined {
     if (!declaration.symbolName) {
         return undefined;
     }
 
-    const aliasData = { originalName: declaration.symbolName!, modulePath: declaration.path };
-    if (!importLookup) {
-        return aliasData;
-    }
-
     const resolved = resolveAliasDeclaration(importLookup, declaration, true);
     const nameValue = resolved ? getNameFromDeclaration(resolved) : undefined;
     if (!nameValue || resolved!.path.length <= 0) {
-        return aliasData;
+        return undefined;
     }
 
-    return { originalName: nameValue, modulePath: resolved!.path };
+    return {
+        originalName: nameValue,
+        modulePath: resolved!.path,
+        kind: getSymbolKind(nameValue, resolved!) ?? SymbolKind.Module,
+    };
 }
 
 // We'll use a somewhat-arbitrary cutoff value here to determine
@@ -76,21 +111,25 @@ export function getIndexAliasData(
 const similarityLimit = 0.5;
 
 export class DocumentSymbolProvider {
-    static addSymbolsForDocument(
+    static getSymbolsForDocument(
         indexResults: IndexResults | undefined,
         parseResults: ParseResults | undefined,
         filePath: string,
         query: string,
-        symbolList: SymbolInformation[],
         token: CancellationToken
-    ) {
+    ): SymbolInformation[] {
+        const symbolList: SymbolInformation[] = [];
+
         if (!indexResults && !parseResults) {
-            return;
+            return symbolList;
         }
 
         const indexSymbolData =
-            indexResults?.symbols ?? DocumentSymbolProvider.indexSymbols(parseResults!, false, token);
+            (indexResults?.symbols as IndexSymbolData[]) ??
+            DocumentSymbolProvider.indexSymbols(parseResults!, { indexingForAutoImportMode: false }, token);
+
         appendWorkspaceSymbolsRecursive(indexSymbolData, filePath, query, '', symbolList, token);
+        return symbolList;
     }
 
     static addHierarchicalSymbolsForDocument(
@@ -104,17 +143,18 @@ export class DocumentSymbolProvider {
         }
 
         const indexSymbolData =
-            indexResults?.symbols ?? DocumentSymbolProvider.indexSymbols(parseResults!, false, token);
+            (indexResults?.symbols as IndexSymbolData[]) ??
+            DocumentSymbolProvider.indexSymbols(parseResults!, { indexingForAutoImportMode: false }, token);
         appendDocumentSymbolsRecursive(indexSymbolData, symbolList, token);
     }
 
     static indexSymbols(
         parseResults: ParseResults,
-        importSymbolsOnly: boolean,
+        options: IndexOptions,
         token: CancellationToken
     ): IndexSymbolData[] {
         const indexSymbolData: IndexSymbolData[] = [];
-        collectSymbolIndexData(parseResults, parseResults.parseTree, importSymbolsOnly, indexSymbolData, token);
+        collectSymbolIndexData(parseResults, parseResults.parseTree, options, indexSymbolData, token);
 
         return indexSymbolData;
     }
@@ -168,7 +208,7 @@ function getSymbolKind(name: string, declaration: Declaration, evaluator?: TypeE
 }
 
 function appendWorkspaceSymbolsRecursive(
-    indexSymbolData: IndexSymbolData[],
+    indexSymbolData: IndexSymbolData[] | undefined,
     filePath: string,
     query: string,
     container: string,
@@ -177,6 +217,10 @@ function appendWorkspaceSymbolsRecursive(
 ) {
     throwIfCancellationRequested(token);
 
+    if (!indexSymbolData) {
+        return;
+    }
+
     for (const symbolData of indexSymbolData) {
         if (symbolData.alias) {
             continue;
@@ -186,7 +230,7 @@ function appendWorkspaceSymbolsRecursive(
         if (similarity >= similarityLimit) {
             const location: Location = {
                 uri: URI.file(filePath).toString(),
-                range: symbolData.selectionRange,
+                range: symbolData.selectionRange!,
             };
 
             const symbolInfo: SymbolInformation = {
@@ -219,12 +263,16 @@ function appendWorkspaceSymbolsRecursive(
 }
 
 function appendDocumentSymbolsRecursive(
-    indexSymbolData: IndexSymbolData[],
+    indexSymbolData: IndexSymbolData[] | undefined,
     symbolList: DocumentSymbol[],
     token: CancellationToken
 ) {
     throwIfCancellationRequested(token);
 
+    if (!indexSymbolData) {
+        return;
+    }
+
     for (const symbolData of indexSymbolData) {
         if (symbolData.alias) {
             continue;
@@ -236,19 +284,41 @@ function appendDocumentSymbolsRecursive(
         const symbolInfo: DocumentSymbol = {
             name: symbolData.name,
             kind: symbolData.kind,
-            range: symbolData.range,
+            range: symbolData.range!,
-            selectionRange: symbolData.selectionRange,
+            selectionRange: symbolData.selectionRange!,
-            children: children,
+            children: children!,
         };
 
         symbolList.push(symbolInfo);
     }
 }
 
+function getAllNameTable(autoImportMode: boolean, root: ModuleNode) {
+    if (!autoImportMode) {
+        // We only care about __all__ in auto import mode.
+        // other cases such as workspace symbols, document symbols, we will collect all symbols
+        // regardless whether it shows up in __all__ or not.
+        return undefined;
+    }
+
+    // If __all__ is defined, we only care ones in the __all__
+    const allNames = AnalyzerNodeInfo.getDunderAllNames(root);
+    if (allNames) {
+        return new Set<string>(allNames);
+    }
+
+    const file = AnalyzerNodeInfo.getFileInfo(root);
+    if (file && file.isStubFile) {
+        return undefined;
+    }
+
+    return new Set<string>();
+}
+
 function collectSymbolIndexData(
     parseResults: ParseResults,
     node: AnalyzerNodeInfo.ScopedNode,
-    autoImportMode: boolean,
+    options: IndexOptions,
     indexSymbolData: IndexSymbolData[],
     token: CancellationToken
 ) {
@@ -259,13 +329,9 @@ function collectSymbolIndexData(
         return;
     }
 
-    // Build __all__ map for regular python file to reduce candidate in autoImportMode.
+    const allNameTable = getAllNameTable(options.indexingForAutoImportMode, parseResults.parseTree);
-    const file = AnalyzerNodeInfo.getFileInfo(parseResults.parseTree);
-    let allNameTable: Set<string> | undefined;
-    if (autoImportMode && !file?.isStubFile) {
-        allNameTable = new Set<string>(AnalyzerNodeInfo.getDunderAllNames(parseResults.parseTree) ?? []);
-    }
 
+    let modulePath: string | undefined = undefined;
     const symbolTable = scope.symbolTable;
     symbolTable.forEach((symbol, name) => {
         if (symbol.isIgnoredForProtocolMatch()) {
@@ -290,11 +356,21 @@ function collectSymbolIndexData(
         }
 
         if (DeclarationType.Alias === declaration.type) {
+            if (!options.indexingForAutoImportMode) {
+                // We don't include import alias for workspace files.
+                return;
+            }
+
             if (declaration.path.length <= 0) {
                 return;
             }
 
-            if (!allNameTable && !includeAliasDeclarationInIndex(declaration)) {
+            const lookup = AnalyzerNodeInfo.getFileInfo(parseResults.parseTree)!.importLookup;
+            modulePath =
+                modulePath ??
+                getLibraryPathWithoutExtension(AnalyzerNodeInfo.getFileInfo(parseResults.parseTree)!.filePath);
+
+            if (!allNameTable && !includeAliasDeclarationInIndex(lookup, modulePath, declaration)) {
                 // For import alias, we only put the alias in the index if it is the form of
                 // "from x import y as z" or the alias is explicitly listed in __all__
                 return;
@@ -304,7 +380,7 @@ function collectSymbolIndexData(
         collectSymbolIndexDataForName(
             parseResults,
             declaration,
-            autoImportMode,
+            options,
             !symbol.isExternallyHidden(),
             name,
             indexSymbolData,
@@ -316,13 +392,13 @@ function collectSymbolIndexData(
 function collectSymbolIndexDataForName(
     parseResults: ParseResults,
     declaration: Declaration,
-    autoImportMode: boolean,
+    options: IndexOptions,
     externallyVisible: boolean,
     name: string,
     indexSymbolData: IndexSymbolData[],
     token: CancellationToken
 ) {
-    if (autoImportMode && !externallyVisible) {
+    if (options.indexingForAutoImportMode && !externallyVisible) {
         return;
     }
 
@@ -336,29 +412,28 @@ function collectSymbolIndexDataForName(
     const children: IndexSymbolData[] = [];
 
     if (declaration.type === DeclarationType.Class || declaration.type === DeclarationType.Function) {
-        if (!autoImportMode) {
+        if (!options.indexingForAutoImportMode) {
-            collectSymbolIndexData(parseResults, declaration.node, false, children, token);
+            collectSymbolIndexData(parseResults, declaration.node, options, children, token);
         }
 
-        const nameRange = convertOffsetsToRange(
+        range = convertOffsetsToRange(
             declaration.node.start,
             declaration.node.name.start + declaration.node.length,
             parseResults.tokenizerOutput.lines
         );
-        range = nameRange;
     }
 
     const data: IndexSymbolData = {
         name,
-        alias:
-            DeclarationType.Alias === declaration.type
-                ? getIndexAliasData(AnalyzerNodeInfo.getFileInfo(parseResults.parseTree)?.importLookup, declaration)
-                : undefined,
         externallyVisible,
         kind: symbolKind,
-        range,
+        alias:
-        selectionRange,
+            DeclarationType.Alias === declaration.type
-        children,
+                ? getIndexAliasData(AnalyzerNodeInfo.getFileInfo(parseResults.parseTree)!.importLookup, declaration)
+                : undefined,
+        range: options.indexingForAutoImportMode ? undefined : range,
+        selectionRange: options.indexingForAutoImportMode ? undefined : selectionRange,
+        children: options.indexingForAutoImportMode ? undefined : children,
     };
 
     indexSymbolData.push(data);

packages/pyright-internal/src/languageService/referencesProvider.ts

@@ -24,15 +24,39 @@ import { TextRange } from '../common/textRange';
 import { ModuleNameNode, NameNode, ParseNode, ParseNodeType } from '../parser/parseNodes';
 import { ParseResults } from '../parser/parser';
 
-export interface ReferencesResult {
+export type ReferenceCallback = (locations: DocumentRange[]) => void;
-    requiresGlobalSearch: boolean;
+
-    nodeAtOffset: ParseNode;
+export class ReferencesResult {
-    symbolName: string;
+    private readonly _locations: DocumentRange[] = [];
-    declarations: Declaration[];
+
-    locations: DocumentRange[];
+    constructor(
+        readonly requiresGlobalSearch: boolean,
+        readonly nodeAtOffset: ParseNode,
+        readonly symbolName: string,
+        readonly declarations: Declaration[],
+        private readonly _reporter?: ReferenceCallback
+    ) {}
+
+    get locations(): readonly DocumentRange[] {
+        return this._locations;
+    }
+
+    addLocations(...locs: DocumentRange[]) {
+        if (locs.length === 0) {
+            return;
+        }
+
+        if (this._reporter) {
+            this._reporter(locs);
+        }
+
+        this._locations.push(...locs);
+    }
 }
 
 class FindReferencesTreeWalker extends ParseTreeWalker {
+    private readonly _locationsFound: DocumentRange[] = [];
+
     constructor(
         private _parseResults: ParseResults,
         private _filePath: string,
@@ -46,6 +70,8 @@ class FindReferencesTreeWalker extends ParseTreeWalker {
 
     findReferences() {
         this.walk(this._parseResults.parseTree);
+
+        return this._locationsFound;
     }
 
     walk(node: ParseNode) {
@@ -74,7 +100,7 @@ class FindReferencesTreeWalker extends ParseTreeWalker {
             if (declarations.some((decl) => this._resultsContainsDeclaration(decl))) {
                 // Is it the same symbol?
                 if (this._includeDeclaration || node !== this._referencesResult.nodeAtOffset) {
-                    this._referencesResult.locations.push({
+                    this._locationsFound.push({
                         path: this._filePath,
                         range: {
                             start: convertOffsetToPosition(node.start, this._parseResults.tokenizerOutput.lines),
@@ -129,6 +155,7 @@ export class ReferencesProvider {
         filePath: string,
         position: Position,
         evaluator: TypeEvaluator,
+        reporter: ReferenceCallback | undefined,
         token: CancellationToken
     ): ReferencesResult | undefined {
         throwIfCancellationRequested(token);
@@ -207,13 +234,7 @@ export class ReferencesProvider {
             return false;
         });
 
-        return {
+        return new ReferencesResult(requiresGlobalSearch, node, node.value, resolvedDeclarations, reporter);
-            requiresGlobalSearch,
-            nodeAtOffset: node,
-            symbolName: node.value,
-            declarations: resolvedDeclarations,
-            locations: [],
-        };
     }
 
     private static _addIfUnique(declarations: Declaration[], itemToAdd: Declaration) {
@@ -242,6 +263,7 @@ export class ReferencesProvider {
             evaluator,
             token
         );
-        refTreeWalker.findReferences();
+
+        referencesResult.addLocations(...refTreeWalker.findReferences());
     }
 }

packages/pyright-internal/src/tests/harness/fourslash/testState.ts

@@ -1006,7 +1006,15 @@ export class TestState {
             const expected = map[name].references;
 
             const position = this.convertOffsetToPosition(fileName, marker.position);
-            const actual = this.program.getReferencesForPosition(fileName, position, true, CancellationToken.None);
+
+            const actual: DocumentRange[] = [];
+            this.program.reportReferencesForPosition(
+                fileName,
+                position,
+                true,
+                (locs) => actual.push(...locs),
+                CancellationToken.None
+            );
 
             assert.equal(actual?.length ?? 0, expected.length);
 

packages/pyright/package.json

@@ -19,7 +19,7 @@
     "scripts": {
         "build": "webpack --mode production --progress",
         "clean": "shx rm -rf ./dist ./out README.md",
-        "prepublishOnly": "npm run clean && shx cp ../../README.md . && npm run build",
+        "prepack": "npm run clean && shx cp ../../README.md . && npm run build",
         "webpack": "webpack --mode development --progress"
     },
     "devDependencies": {