Improved startup performance for large projects by lazily computing the module name of each source file. (#10798)

Author: erictraut

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

@@ -336,21 +336,19 @@ export class Program {
 
     addTrackedFile(fileUri: Uri, isThirdPartyImport = false, isInPyTypedPackage = false): SourceFile {
         let sourceFileInfo = this.getSourceFileInfo(fileUri);
-        const moduleImportInfo = this._getModuleImportInfoForFile(fileUri);
-        const importName = moduleImportInfo.moduleName;
 
         if (sourceFileInfo) {
             // The module name may have changed based on updates to the
-            // search paths, so update it here.
-            sourceFileInfo.sourceFile.setModuleName(importName);
+            // search paths. Clear any cached module name so it is recomputed.
+            sourceFileInfo.sourceFile.clearCachedModuleName();
             sourceFileInfo.isTracked = true;
             return sourceFileInfo.sourceFile;
         }
 
         const sourceFile = this._sourceFileFactory.createSourceFile(
             this.serviceProvider,
             fileUri,
-            importName,
+            (uri) => this._getModuleName(uri),
             isThirdPartyImport,
             isInPyTypedPackage,
             this._editModeTracker,
@@ -378,7 +376,7 @@ export class Program {
             const sourceFile = this._sourceFileFactory.createSourceFile(
                 this.serviceProvider,
                 fileUri,
-                moduleImportInfo.moduleName,
+                (uri) => this._getModuleName(uri),
                 /* isThirdPartyImport */ false,
                 moduleImportInfo.isThirdPartyPyTypedPresent,
                 this._editModeTracker,
@@ -1506,11 +1504,10 @@ export class Program {
                 // of the program.
                 let importedFileInfo = this.getSourceFileInfo(importInfo.path);
                 if (!importedFileInfo) {
-                    const moduleImportInfo = this._getModuleImportInfoForFile(importInfo.path);
                     const sourceFile = this._sourceFileFactory.createSourceFile(
                         this.serviceProvider,
                         importInfo.path,
-                        moduleImportInfo.moduleName,
+                        (uri) => this._getModuleName(uri),
                         importInfo.isThirdPartyImport,
                         importInfo.isPyTypedPresent,
                         this._editModeTracker,
@@ -1582,6 +1579,11 @@ export class Program {
         this._sourceFileMap.set(fileUri.key, fileInfo);
     }
 
+    private _getModuleName(fileUri: Uri): string {
+        const moduleInfo = this._getModuleImportInfoForFile(fileUri);
+        return moduleInfo.moduleName;
+    }
+
     private _getModuleImportInfoForFile(fileUri: Uri) {
         // We allow illegal module names (e.g. names that include "-" in them)
         // because we want a unique name for each module even if it cannot be
@@ -1622,11 +1624,10 @@ export class Program {
     }
 
     private _createInterimFileInfo(fileUri: Uri) {
-        const moduleImportInfo = this._getModuleImportInfoForFile(fileUri);
         const sourceFile = this._sourceFileFactory.createSourceFile(
             this.serviceProvider,
             fileUri,
-            moduleImportInfo.moduleName,
+            (uri) => this._getModuleName(uri),
             /* isThirdPartyImport */ false,
             /* isInPyTypedPackage */ false,
             this._editModeTracker,

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

@@ -9,13 +9,13 @@ import { ConsoleInterface } from '../common/console';
 import { LogTracker } from '../common/logTracker';
 import { ServiceProvider } from '../common/serviceProvider';
 import { Uri } from '../common/uri/uri';
-import { SourceFileEditMode, IPythonMode, SourceFile } from './sourceFile';
+import { IPythonMode, SourceFile, SourceFileEditMode } from './sourceFile';
 
 export interface ISourceFileFactory {
     createSourceFile(
         serviceProvider: ServiceProvider,
         fileUri: Uri,
-        moduleName: string,
+        moduleNameGetter: (file: Uri) => string,
         isThirdPartyImport: boolean,
         isThirdPartyPyTypedPresent: boolean,
         editMode: SourceFileEditMode,

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

@@ -206,8 +206,11 @@ export class SourceFile {
     // identify this file.
     private readonly _fileId: string;
 
+    // Getter to lazily compute the module name from the file URI.
+    private _moduleNameGetter: (file: Uri) => string;
+
     // Period-delimited import path for the module.
-    private _moduleName: string;
+    private _cachedModuleName: string | undefined;
 
     // True if file is a type-hint (.pyi) file versus a python
     // (.py) file.
@@ -257,7 +260,7 @@ export class SourceFile {
     constructor(
         readonly serviceProvider: ServiceProvider,
         uri: Uri,
-        moduleName: string,
+        moduleNameGetter: (file: Uri) => string,
         isThirdPartyImport: boolean,
         isThirdPartyPyTypedPresent: boolean,
         editMode: SourceFileEditMode,
@@ -272,7 +275,7 @@ export class SourceFile {
         this._editMode = editMode;
         this._uri = uri;
         this._fileId = this._makeFileId(uri);
-        this._moduleName = moduleName;
+        this._moduleNameGetter = moduleNameGetter;
         this._isStubFile = uri.hasExtension('.pyi');
         this._isThirdPartyImport = isThirdPartyImport;
         this._isThirdPartyPyTypedPresent = isThirdPartyPyTypedPresent;
@@ -319,16 +322,17 @@ export class SourceFile {
     }
 
     getModuleName(): string {
-        if (this._moduleName) {
-            return this._moduleName;
+        if (!this._cachedModuleName) {
+            // Call the module name getter. If it returns '' (which can happen if the file is not part
+            // of the project), fall back to the file name.)
+            return this._moduleNameGetter(this._uri) || stripFileExtension(this._uri.fileName);
         }
 
-        // Synthesize a module name using the file path.
-        return stripFileExtension(this._uri.fileName);
+        return this._cachedModuleName;
     }
 
-    setModuleName(name: string) {
-        this._moduleName = name;
+    clearCachedModuleName() {
+        this._cachedModuleName = undefined;
     }
 
     getDiagnosticVersion(): number {

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

@@ -115,7 +115,7 @@ const DefaultSourceFileFactory: ISourceFileFactory = {
     createSourceFile(
         serviceProvider: ServiceProvider,
         fileUri: Uri,
-        moduleName: string,
+        moduleNameGetter: (file: Uri) => string,
         isThirdPartyImport: boolean,
         isThirdPartyPyTypedPresent: boolean,
         editMode: SourceFileEditMode,
@@ -126,7 +126,7 @@ const DefaultSourceFileFactory: ISourceFileFactory = {
         return new SourceFile(
             serviceProvider,
             fileUri,
-            moduleName,
+            moduleNameGetter,
             isThirdPartyImport,
             isThirdPartyPyTypedPresent,
             editMode,

packages/pyright-internal/src/tests/sourceFile.test.ts

@@ -23,7 +23,7 @@ test('Empty', () => {
     const tempFile = new RealTempFile();
     const fs = createFromRealFileSystem(tempFile);
     const serviceProvider = createServiceProvider(tempFile, fs);
-    const sourceFile = new SourceFile(serviceProvider, Uri.file(filePath, serviceProvider), '', false, false, {
+    const sourceFile = new SourceFile(serviceProvider, Uri.file(filePath, serviceProvider), () => '', false, false, {
         isEditMode: false,
     });
     const configOptions = new ConfigOptions(Uri.file(process.cwd(), serviceProvider));