Implemented a number of performance and memory optimizations (#10816)

erictraut · web-flow · commit 572bc31e9b0a · 2025-08-17T22:26:02.000-07:00
* Implemented a number of performance and memory optimizations:
* Made implicitImports and filteredImplicitImports fields of ImportResult optional to eliminate the need to create empty maps.
* Avoided creating a new diagnostic rule set for every source file unless it overrides the default rule set.
* Limited memoization of file and web URIs; we were previously allocating and caching an unlimited number of URIs which consumed 100's of Mb of space on large project.
* Changed tokenizer to forcibly clone the value of a string literal so we don't retain a reference to the entire file's text contents.

* Fixed regression caught by mypy_primer.
diff --git a/packages/pyright-internal/src/analyzer/binder.ts b/packages/pyright-internal/src/analyzer/binder.ts
@@ -400,7 +400,7 @@ export class Binder extends ParseTreeWalker {
             if (importResult.isNamespacePackage && node.parent?.nodeType === ParseNodeType.ImportFrom) {
                 if (
                     node.parent.d.imports.every((importAs) => {
-                        const implicitImport = importResult.filteredImplicitImports.get(importAs.d.name.d.value);
+                        const implicitImport = importResult.filteredImplicitImports?.get(importAs.d.name.d.value);
                         return !!implicitImport?.pyTypedInfo;
                     })
                 ) {
@@ -4168,7 +4168,7 @@ export class Binder extends ParseTreeWalker {
     }
 
     private _addImplicitImportsToLoaderActions(importResult: ImportResult, loaderActions: ModuleLoaderActions) {
-        importResult.filteredImplicitImports.forEach((implicitImport) => {
+        importResult.filteredImplicitImports?.forEach((implicitImport) => {
             const existingLoaderAction = loaderActions.implicitImports
                 ? loaderActions.implicitImports.get(implicitImport.name)
                 : undefined;
diff --git a/packages/pyright-internal/src/analyzer/commentUtils.ts b/packages/pyright-internal/src/analyzer/commentUtils.ts
@@ -45,9 +45,11 @@ export function getFileLevelDirectives(
     diagnostics: CommentDiagnostic[]
 ): DiagnosticRuleSet {
     let ruleSet = cloneDiagnosticRuleSet(defaultRuleSet);
+    let isModified = false;
 
     if (useStrict) {
         _applyStrictRules(ruleSet);
+        isModified = true;
     }
 
     for (let i = 0; i < tokens.count; i++) {
@@ -65,11 +67,13 @@ export function getFileLevelDirectives(
                 };
 
                 ruleSet = _parsePyrightComment(value, textRange, isCommentOnOwnLine, ruleSet, diagnostics);
+                isModified = true;
             }
         }
     }
 
-    return ruleSet;
+    // If we didn't make any modifications, use the default rule set to save memory.
+    return isModified ? ruleSet : defaultRuleSet;
 }
 
 function _applyStrictRules(ruleSet: DiagnosticRuleSet) {
diff --git a/packages/pyright-internal/src/analyzer/importResolver.ts b/packages/pyright-internal/src/analyzer/importResolver.ts
@@ -442,8 +442,6 @@ export class ImportResolver {
             this._cachedPythonSearchPaths = { paths: Array.from(new Set(paths)), failureInfo: info };
         }
 
-        // Make sure we cache the logs as well so we can find out why search path failed.
-        importFailureInfo.push(...this._cachedPythonSearchPaths.failureInfo);
         return this._cachedPythonSearchPaths.paths;
     }
 
@@ -791,15 +789,15 @@ export class ImportResolver {
     ): ImportResult {
         if (importedSymbols === undefined) {
             const newImportResult = Object.assign({}, importResult);
-            newImportResult.filteredImplicitImports = new Map<string, ImplicitImport>();
+            newImportResult.filteredImplicitImports = undefined;
             return newImportResult;
         }
 
-        if (importedSymbols.size === 0) {
+        if (importedSymbols === undefined || importedSymbols.size === 0) {
             return importResult;
         }
 
-        if (importResult.implicitImports.size === 0) {
+        if (importResult.implicitImports === undefined || importResult.implicitImports.size === 0) {
             return importResult;
         }
 
@@ -823,7 +821,7 @@ export class ImportResolver {
         importingModuleName: string,
         dirPath: Uri,
         exclusions: Uri[]
-    ): Map<string, ImplicitImport> {
+    ): Map<string, ImplicitImport> | undefined {
         const implicitImportMap = new Map<string, ImplicitImport>();
 
         // Enumerate all of the files and directories in the path, expanding links.
@@ -911,7 +909,7 @@ export class ImportResolver {
             }
         }
 
-        return implicitImportMap;
+        return implicitImportMap.size > 0 ? implicitImportMap : undefined;
     }
 
     private _resolveImportStrict(
@@ -935,8 +933,8 @@ export class ImportResolver {
             importType: ImportType.Local,
             isStubFile: false,
             isNativeLib: false,
-            implicitImports: new Map<string, ImplicitImport>(),
-            filteredImplicitImports: new Map<string, ImplicitImport>(),
+            implicitImports: undefined,
+            filteredImplicitImports: undefined,
             nonStubImportResult: undefined,
         };
 
@@ -1352,7 +1350,7 @@ export class ImportResolver {
         let isStubPackage = false;
         let isStubFile = false;
         let isNativeLib = false;
-        let implicitImports = new Map<string, ImplicitImport>();
+        let implicitImports: Map<string, ImplicitImport> | undefined;
         let packageDirectory: Uri | undefined;
         let pyTypedInfo: PyTypedInfo | undefined;
 
@@ -1553,13 +1551,13 @@ export class ImportResolver {
     // are all satisfied by submodules (as listed in the implicit imports).
     private _isNamespacePackageResolved(
         moduleDescriptor: ImportedModuleDescriptor,
-        implicitImports: Map<string, ImplicitImport>
+        implicitImports: Map<string, ImplicitImport> | undefined
     ) {
         if (moduleDescriptor.importedSymbols) {
-            if (!Array.from(moduleDescriptor.importedSymbols.keys()).some((symbol) => implicitImports.has(symbol))) {
+            if (!Array.from(moduleDescriptor.importedSymbols.keys()).some((symbol) => implicitImports?.has(symbol))) {
                 return false;
             }
-        } else if (implicitImports.size === 0) {
+        } else if (!implicitImports || implicitImports.size === 0) {
             return false;
         }
         return true;
@@ -2377,8 +2375,8 @@ export class ImportResolver {
                 importType: ImportType.Local,
                 isStubFile: false,
                 isNativeLib: false,
-                implicitImports: [],
-                filteredImplicitImports: [],
+                implicitImports: undefined,
+                filteredImplicitImports: undefined,
                 nonStubImportResult: undefined,
             };
         }
diff --git a/packages/pyright-internal/src/analyzer/importResult.ts b/packages/pyright-internal/src/analyzer/importResult.ts
@@ -88,12 +88,12 @@ export interface ImportResult {
     // List of files within the final resolved path that are implicitly
     // imported as part of the package - used for both traditional and
     // namespace packages.
-    implicitImports: Map<string, ImplicitImport>;
+    implicitImports?: Map<string, ImplicitImport>;
 
     // Implicit imports that have been filtered to include only
     // those symbols that are explicitly imported in a "from x import y"
     // statement.
-    filteredImplicitImports: Map<string, ImplicitImport>;
+    filteredImplicitImports?: Map<string, ImplicitImport>;
 
     // If resolved from a type hint (.pyi), then store the import result
     // from .py here.
diff --git a/packages/pyright-internal/src/analyzer/importStatementUtils.ts b/packages/pyright-internal/src/analyzer/importStatementUtils.ts
@@ -678,10 +678,12 @@ function _processImportFromNode(
     if (includeImplicitImports && importResult) {
         localImports.implicitImports = localImports.implicitImports ?? new Map<string, ImportFromAsNode>();
 
-        for (const implicitImport of importResult.implicitImports.values()) {
-            const importFromAs = node.d.imports.find((i) => i.d.name.d.value === implicitImport.name);
-            if (importFromAs) {
-                localImports.implicitImports.set(implicitImport.uri.key, importFromAs);
+        if (importResult.implicitImports) {
+            for (const implicitImport of importResult.implicitImports.values()) {
+                const importFromAs = node.d.imports.find((i) => i.d.name.d.value === implicitImport.name);
+                if (importFromAs) {
+                    localImports.implicitImports.set(implicitImport.uri.key, importFromAs);
+                }
             }
         }
     }
diff --git a/packages/pyright-internal/src/analyzer/program.ts b/packages/pyright-internal/src/analyzer/program.ts
@@ -1339,6 +1339,7 @@ export class Program {
                 thirdPartyImportAllowed = true;
             } else if (
                 importResult.isNamespacePackage &&
+                importResult.filteredImplicitImports &&
                 Array.from(importResult.filteredImplicitImports.values()).some(
                     (implicitImport) => !!implicitImport.pyTypedInfo
                 )
@@ -1434,7 +1435,7 @@ export class Program {
                     }
                 }
 
-                importResult.filteredImplicitImports.forEach((implicitImport) => {
+                importResult.filteredImplicitImports?.forEach((implicitImport) => {
                     if (this._isImportAllowed(sourceFileInfo, importResult, implicitImport.isStubFile)) {
                         if (!implicitImport.isNativeLib) {
                             const thirdPartyTypeInfo = getThirdPartyImportInfo(importResult);
diff --git a/packages/pyright-internal/src/analyzer/service.ts b/packages/pyright-internal/src/analyzer/service.ts
@@ -1407,7 +1407,7 @@ export class AnalyzerService {
                 }
 
                 // Add the implicit import paths.
-                importResult.filteredImplicitImports.forEach((implicitImport) => {
+                importResult.filteredImplicitImports?.forEach((implicitImport) => {
                     if (ImportResolver.isSupportedImportSourceFile(implicitImport.uri)) {
                         filesToImport.push(implicitImport.uri);
                     }
diff --git a/packages/pyright-internal/src/analyzer/sourceFile.ts b/packages/pyright-internal/src/analyzer/sourceFile.ts
@@ -455,6 +455,7 @@ export class SourceFile {
         this._writableData.parsedFileContents = undefined;
         this._writableData.moduleSymbolTable = undefined;
         this._writableData.isBindingNeeded = true;
+        this._writableData.imports = [];
     }
 
     markDirty(): void {
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -643,7 +643,6 @@ export function createTypeEvaluator(
     wrapWithLogger: LogWrapper
 ): TypeEvaluator {
     const symbolResolutionStack: SymbolResolutionStackEntry[] = [];
-    const asymmetricAccessorAssignmentCache = new Set<number>();
     const speculativeTypeTracker = new SpeculativeTypeTracker();
     const suppressedNodeStack: SuppressedNodeStackEntry[] = [];
     const assignClassToSelfStack: AssignClassToSelfInfo[] = [];
@@ -653,6 +652,7 @@ export function createTypeEvaluator(
     let typeCache = new Map<number, TypeCacheEntry>();
     let effectiveTypeCache = new Map<number, Map<string, EffectiveTypeResult>>();
     let expectedTypeCache = new Map<number, Type>();
+    let asymmetricAccessorAssignmentCache = new Set<number>();
     let deferredClassCompletions: DeferredClassCompletion[] = [];
     let cancellationToken: CancellationToken | undefined;
     let printExpressionSpaceCount = 0;
@@ -708,6 +708,7 @@ export function createTypeEvaluator(
         typeCache = new Map<number, TypeCacheEntry>();
         effectiveTypeCache = new Map<number, Map<string, EffectiveTypeResult>>();
         expectedTypeCache = new Map<number, Type>();
+        asymmetricAccessorAssignmentCache = new Set<number>();
     }
 
     function readTypeCacheEntry(node: ParseNode) {
diff --git a/packages/pyright-internal/src/common/fullAccessHost.ts b/packages/pyright-internal/src/common/fullAccessHost.ts
@@ -15,12 +15,12 @@ import { PythonPlatform } from './configOptions';
 import { assertNever } from './debug';
 import { HostKind, NoAccessHost, ScriptOutput } from './host';
 import { getAnyExtensionFromPath, normalizePath } from './pathUtils';
+import { terminateChild } from './processUtils';
 import { PythonVersion } from './pythonVersion';
 import { ServiceKeys } from './serviceKeys';
 import { ServiceProvider } from './serviceProvider';
 import { Uri } from './uri/uri';
 import { isDirectory } from './uri/uriUtils';
-import { terminateChild } from './processUtils';
 
 // preventLocalImports removes the working directory from sys.path.
 // The -c flag adds it automatically, which can allow some stdlib
diff --git a/packages/pyright-internal/src/common/uri/memoization.ts b/packages/pyright-internal/src/common/uri/memoization.ts
@@ -6,9 +6,12 @@
  * Decorators used to memoize the result of a function call.
  */
 
-// Cache for static method results.
+// Cache for static method results with LRU eviction.
 const staticCache = new Map<string, any>();
 
+// Max number of static method values that are cached.
+const maxStaticCacheEntries = 256;
+
 // Caches the results of a getter property.
 export function cacheProperty() {
     return function (target: any, functionName: string, descriptor: PropertyDescriptor) {
@@ -56,14 +59,26 @@ export function cacheStaticFunc() {
         const originalMethod = descriptor.value;
         descriptor.value = function (...args: any) {
             const key = `${functionName}+${args?.map((a: any) => a?.toString()).join(',')}`;
-            let cachedResult: any;
-            if (!staticCache.has(key)) {
-                cachedResult = originalMethod.apply(this, args);
-                staticCache.set(key, cachedResult);
-            } else {
-                cachedResult = staticCache.get(key);
+            if (staticCache.has(key)) {
+                // Promote to most-recently used by re-inserting.
+                const value = staticCache.get(key);
+                staticCache.delete(key);
+                staticCache.set(key, value);
+                return value;
+            }
+
+            // Miss: compute and insert, evict LRU if over capacity.
+            const result = originalMethod.apply(this, args);
+
+            if (staticCache.size >= maxStaticCacheEntries) {
+                // Remove least-recently used (the first key in insertion order).
+                const lruKey = staticCache.keys().next().value as string | undefined;
+                if (lruKey !== undefined) {
+                    staticCache.delete(lruKey);
+                }
             }
-            return cachedResult;
+            staticCache.set(key, result);
+            return result;
         };
         return descriptor;
     };
diff --git a/packages/pyright-internal/src/common/uri/webUri.ts b/packages/pyright-internal/src/common/uri/webUri.ts
@@ -11,11 +11,11 @@
  * - vscode-vfs://github.com/microsoft/debugpy/debugpy/launcher/debugAdapter.py
  */
 
+import { URI } from 'vscode-uri';
 import { getRootLength, hasTrailingDirectorySeparator, resolvePaths } from '../pathUtils';
 import { BaseUri, JsonObjType } from './baseUri';
 import { cacheMethodWithNoArgs, cacheProperty, cacheStaticFunc } from './memoization';
 import { Uri } from './uri';
-import { URI } from 'vscode-uri';
 
 export class WebUri extends BaseUri {
     private constructor(
diff --git a/packages/pyright-internal/src/languageService/completionProvider.ts b/packages/pyright-internal/src/languageService/completionProvider.ts
@@ -2826,7 +2826,7 @@ export class CompletionProvider {
         priorWord: string,
         completionMap: CompletionMap
     ) {
-        importInfo.implicitImports.forEach((implImport) => {
+        importInfo.implicitImports?.forEach((implImport) => {
             if (!importFromNode.d.imports.find((imp) => imp.d.name.d.value === implImport.name)) {
                 this.addNameToCompletions(implImport.name, CompletionItemKind.Module, priorWord, completionMap, {
                     moduleUri: implImport.uri,
diff --git a/packages/pyright-internal/src/parser/tokenizer.ts b/packages/pyright-internal/src/parser/tokenizer.ts
@@ -1616,7 +1616,7 @@ export class Tokenizer {
         let isInNamedUnicodeEscape = false;
         const start = this._cs.position;
         let escapedValueLength = 0;
-        const getEscapedValue = () => this._cs.getText().slice(start, start + escapedValueLength);
+        const getEscapedValue = () => cloneStr(this._cs.getText().slice(start, start + escapedValueLength));
 
         while (true) {
             if (this._cs.isEndOfStream()) {

Original file line number	Diff line number	Diff line change
`@@ -45,9 +45,11 @@ export function getFileLevelDirectives(`
`45`	`45`	`diagnostics: CommentDiagnostic[]`
`46`	`46`	`): DiagnosticRuleSet {`
`47`	`47`	`let ruleSet = cloneDiagnosticRuleSet(defaultRuleSet);`
	`48`	`+ let isModified = false;`
`48`	`49`
`49`	`50`	`if (useStrict) {`
`50`	`51`	`_applyStrictRules(ruleSet);`
	`52`	`+ isModified = true;`
`51`	`53`	`}`
`52`	`54`
`53`	`55`	`for (let i = 0; i < tokens.count; i++) {`
`@@ -65,11 +67,13 @@ export function getFileLevelDirectives(`
`65`	`67`	`};`
`66`	`68`
`67`	`69`	`ruleSet = _parsePyrightComment(value, textRange, isCommentOnOwnLine, ruleSet, diagnostics);`
	`70`	`+ isModified = true;`
`68`	`71`	`}`
`69`	`72`	`}`
`70`	`73`	`}`
`71`	`74`
`72`		`- return ruleSet;`
	`75`	`+ // If we didn't make any modifications, use the default rule set to save memory.`
	`76`	`+ return isModified ? ruleSet : defaultRuleSet;`
`73`	`77`	`}`
`74`	`78`
`75`	`79`	`function _applyStrictRules(ruleSet: DiagnosticRuleSet) {`
Original file line number	Diff line number	Diff line change
`@@ -678,10 +678,12 @@ function _processImportFromNode(`
`678`	`678`	`if (includeImplicitImports && importResult) {`
`679`	`679`	`localImports.implicitImports = localImports.implicitImports ?? new Map<string, ImportFromAsNode>();`
`680`	`680`
`681`		`- for (const implicitImport of importResult.implicitImports.values()) {`
`682`		`- const importFromAs = node.d.imports.find((i) => i.d.name.d.value === implicitImport.name);`
`683`		`- if (importFromAs) {`
`684`		`- localImports.implicitImports.set(implicitImport.uri.key, importFromAs);`
	`681`	`+ if (importResult.implicitImports) {`
	`682`	`+ for (const implicitImport of importResult.implicitImports.values()) {`
	`683`	`+ const importFromAs = node.d.imports.find((i) => i.d.name.d.value === implicitImport.name);`
	`684`	`+ if (importFromAs) {`
	`685`	`+ localImports.implicitImports.set(implicitImport.uri.key, importFromAs);`
	`686`	`+ }`
`685`	`687`	`}`
`686`	`688`	`}`
`687`	`689`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1339,6 +1339,7 @@ export class Program {`
`1339`	`1339`	`thirdPartyImportAllowed = true;`
`1340`	`1340`	`} else if (`
`1341`	`1341`	`importResult.isNamespacePackage &&`
	`1342`	`+ importResult.filteredImplicitImports &&`
`1342`	`1343`	`Array.from(importResult.filteredImplicitImports.values()).some(`
`1343`	`1344`	`(implicitImport) => !!implicitImport.pyTypedInfo`
`1344`	`1345`	`)`
`@@ -1434,7 +1435,7 @@ export class Program {`
`1434`	`1435`	`}`
`1435`	`1436`	`}`
`1436`	`1437`
`1437`		`- importResult.filteredImplicitImports.forEach((implicitImport) => {`
	`1438`	`+ importResult.filteredImplicitImports?.forEach((implicitImport) => {`
`1438`	`1439`	`if (this._isImportAllowed(sourceFileInfo, importResult, implicitImport.isStubFile)) {`
`1439`	`1440`	`if (!implicitImport.isNativeLib) {`
`1440`	`1441`	`const thirdPartyTypeInfo = getThirdPartyImportInfo(importResult);`
Original file line number	Diff line number	Diff line change
`@@ -1407,7 +1407,7 @@ export class AnalyzerService {`
`1407`	`1407`	`}`
`1408`	`1408`
`1409`	`1409`	`// Add the implicit import paths.`
`1410`		`- importResult.filteredImplicitImports.forEach((implicitImport) => {`
	`1410`	`+ importResult.filteredImplicitImports?.forEach((implicitImport) => {`
`1411`	`1411`	`if (ImportResolver.isSupportedImportSourceFile(implicitImport.uri)) {`
`1412`	`1412`	`filesToImport.push(implicitImport.uri);`
`1413`	`1413`	`}`
Original file line number	Diff line number	Diff line change
`@@ -455,6 +455,7 @@ export class SourceFile {`
`455`	`455`	`this._writableData.parsedFileContents = undefined;`
`456`	`456`	`this._writableData.moduleSymbolTable = undefined;`
`457`	`457`	`this._writableData.isBindingNeeded = true;`
	`458`	`+ this._writableData.imports = [];`
`458`	`459`	`}`
`459`	`460`
`460`	`461`	`markDirty(): void {`