Idea for copy on write

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

@@ -24,6 +24,7 @@ export class BackgroundAnalysisProgram {
     private _program: Program;
     private _disposed = false;
     private _onAnalysisCompletion: AnalysisCompleteCallback | undefined;
+    private _preEditAnalysis: BackgroundAnalysisBase | undefined;
 
     constructor(
         private _console: ConsoleInterface,
@@ -242,6 +243,21 @@ export class BackgroundAnalysisProgram {
         this._backgroundAnalysis?.shutdown();
     }
 
+    enterEditMode() {
+        // Turn off analysis while in edit mode.
+        this._preEditAnalysis = this._backgroundAnalysis;
+        this._backgroundAnalysis = undefined;
+
+        // Forward this request to the program.
+        this._program.enterEditMode();
+    }
+
+    exitEditMode() {
+        this._backgroundAnalysis = this._preEditAnalysis;
+        this._preEditAnalysis = undefined;
+        return this._program.exitEditMode();
+    }
+
     protected getIndices(): Indices | undefined {
         return undefined;
     }

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

@@ -11,6 +11,7 @@
 import { CancellationToken, CompletionItem, DocumentSymbol } from 'vscode-languageserver';
 import { CompletionList } from 'vscode-languageserver-types';
 
+import { TextDocument } from 'vscode-languageserver-textdocument';
 import { Commands } from '../commands/commands';
 import { OperationCanceledException, throwIfCancellationRequested } from '../common/cancellationUtils';
 import { appendArray, arrayEquals } from '../common/collectionUtils';
@@ -38,17 +39,17 @@ import {
 import { convertPositionToOffset, convertRangeToTextRange, convertTextRangeToRange } from '../common/positionUtils';
 import { computeCompletionSimilarity } from '../common/stringUtils';
 import { TextEditTracker } from '../common/textEditTracker';
-import { doRangesIntersect, getEmptyRange, Position, Range, TextRange } from '../common/textRange';
+import { Position, Range, TextRange, doRangesIntersect, getEmptyRange } from '../common/textRange';
 import { TextRangeCollection } from '../common/textRangeCollection';
 import { Duration, timingStats } from '../common/timing';
 import { applyTextEditsToString } from '../common/workspaceEditUtils';
 import {
-    AutoImporter,
     AutoImportOptions,
     AutoImportResult,
-    buildModuleSymbolsMap,
+    AutoImporter,
     ImportFormat,
     ModuleSymbolMap,
+    buildModuleSymbolsMap,
 } from '../languageService/autoImporter';
 import {
     AbbreviationMap,
@@ -84,15 +85,15 @@ import { Scope } from './scope';
 import { getScopeForNode } from './scopeUtils';
 import { IPythonMode, SourceFile } from './sourceFile';
 import { collectImportedByFiles, isUserCode } from './sourceFileInfoUtils';
-import { isStubFile, SourceMapper } from './sourceMapper';
+import { SourceMapper, isStubFile } from './sourceMapper';
 import { Symbol } from './symbol';
 import { isPrivateOrProtectedName } from './symbolNameUtils';
 import { createTracePrinter } from './tracePrinter';
 import { PrintTypeOptions, TypeEvaluator } from './typeEvaluatorTypes';
 import { createTypeEvaluatorWithTracker } from './typeEvaluatorWithTracker';
 import { PrintTypeFlags } from './typePrinter';
-import { Type } from './types';
 import { TypeStubWriter } from './typeStubWriter';
+import { Type } from './types';
 
 const _maxImportDepth = 256;
 
@@ -187,6 +188,7 @@ export class Program {
     private _cacheManager: CacheManager;
     private _id: number;
     private static _nextId = 0;
+    private _isEditMode = false;
 
     constructor(
         initialImportResolver: ImportResolver,
@@ -240,6 +242,53 @@ export class Program {
         this._cacheManager.unregisterCacheOwner(this);
     }
 
+    enterEditMode() {
+        // Keep track of edit mode so we can apply it to new source files.
+        this._isEditMode = true;
+
+        // Tell all source files we're in edit mode.
+        this._sourceFileList.forEach((sourceFile) => {
+            sourceFile.sourceFile.enterEditMode();
+        });
+    }
+
+    exitEditMode() {
+        // Tell all source files we're no longer in edit mode. Gather
+        // up all of their edits.
+        const edits: FileEditAction[] = [];
+        this._sourceFileList.forEach((sourceFile) => {
+            const newContents = sourceFile.sourceFile.exitEditMode();
+            if (newContents) {
+                // This means this source file was modified. We need to recompute its imports after
+                // we put it back to how it was.
+                this._updateSourceFileImports(sourceFile, this.configOptions);
+
+                // Create a text document so we can compute the edits.
+                const textDocument = TextDocument.create(
+                    sourceFile.sourceFile.getFilePath(),
+                    'python',
+                    1,
+                    sourceFile.sourceFile.getFileContent() || ''
+                );
+
+                // Add an edit action to the list.
+                edits.push({
+                    filePath: sourceFile.sourceFile.getFilePath(),
+                    range: {
+                        start: { line: 0, character: 0 },
+                        end: { line: textDocument.lineCount, character: 0 },
+                    },
+                    replacementText: newContents,
+                });
+            }
+        });
+
+        // Stop applying edit mode to new source files.
+        this._isEditMode = false;
+
+        return edits;
+    }
+
     setConfigOptions(configOptions: ConfigOptions) {
         this._configOptions = configOptions;
         this._importResolver.setConfigOptions(configOptions);
@@ -331,6 +380,7 @@ export class Program {
             importName,
             isThirdPartyImport,
             isInPyTypedPackage,
+            this._isEditMode,
             this._console,
             this._logTracker
         );
@@ -361,12 +411,12 @@ export class Program {
                 importName,
                 /* isThirdPartyImport */ false,
                 /* isInPyTypedPackage */ false,
+                this._isEditMode,
                 this._console,
                 this._logTracker,
                 options?.realFilePath,
                 options?.ipythonMode ?? IPythonMode.None
             );
-
             const chainedFilePath = options?.chainedFilePath;
             sourceFileInfo = {
                 sourceFile,
@@ -2501,6 +2551,7 @@ export class Program {
                         importName,
                         importInfo.isThirdPartyImport,
                         importInfo.isPyTypedPresent,
+                        this._isEditMode,
                         this._console,
                         this._logTracker
                     );
@@ -2644,6 +2695,7 @@ export class Program {
             importName,
             /* isThirdPartyImport */ false,
             /* isInPyTypedPackage */ false,
+            this._isEditMode,
             this._console,
             this._logTracker
         );

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

@@ -23,7 +23,7 @@ import { CommandLineOptions } from '../common/commandLineOptions';
 import { ConfigOptions, matchFileSpecs } from '../common/configOptions';
 import { ConsoleInterface, LogLevel, StandardConsole, log } from '../common/console';
 import { Diagnostic } from '../common/diagnostic';
-import { FileEditActions } from '../common/editAction';
+import { FileEditAction, FileEditActions } from '../common/editAction';
 import { Extensions, ProgramView } from '../common/extensibility';
 import { FileSystem, FileWatcher, FileWatcherEventType, ignoredWatchEventFunction } from '../common/fileSystem';
 import { Host, HostFactory, NoAccessHost } from '../common/host';
@@ -222,6 +222,22 @@ export class AnalyzerService {
         return service;
     }
 
+    useEditMode(callback: () => void, token: CancellationToken) {
+        let edits: FileEditAction[] = [];
+        const disposable = token.onCancellationRequested(() => {
+            edits = [];
+            this._backgroundAnalysisProgram.exitEditMode();
+        });
+        this._backgroundAnalysisProgram.enterEditMode();
+        try {
+            callback();
+        } finally {
+            disposable.dispose();
+            edits = this._backgroundAnalysisProgram.exitEditMode();
+        }
+        return edits;
+    }
+
     dispose() {
         if (!this._disposed) {
             // Make sure we dispose program, otherwise, entire program