fix baseline when running in multithreaded/background analysis mode

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

@@ -906,6 +906,7 @@ export class Program {
                         fileUri: sourceFileInfo.sourceFile.getUri(),
                         version: sourceFileInfo.sourceFile.getClientVersion(),
                         diagnostics,
+                        reason: 'analysis',
                     });
 
                     // Update the cached diagnosticsVersion so we can determine
@@ -923,6 +924,7 @@ export class Program {
                     fileUri: sourceFileInfo.sourceFile.getUri(),
                     version: sourceFileInfo.sourceFile.getClientVersion(),
                     diagnostics: [],
+                    reason: 'tracking',
                 });
                 sourceFileInfo.diagnosticsVersion = undefined;
             }
@@ -1087,6 +1089,7 @@ export class Program {
                         fileUri: fileInfo.sourceFile.getUri(),
                         version: fileInfo.sourceFile.getClientVersion(),
                         diagnostics: [],
+                        reason: 'tracking',
                     });
                 }
 
@@ -1115,6 +1118,7 @@ export class Program {
                                     fileUri: importedFile.sourceFile.getUri(),
                                     version: importedFile.sourceFile.getClientVersion(),
                                     diagnostics: [],
+                                    reason: 'tracking',
                                 });
                             }
 
@@ -1138,6 +1142,7 @@ export class Program {
                         fileUri: fileInfo.sourceFile.getUri(),
                         version: fileInfo.sourceFile.getClientVersion(),
                         diagnostics: [],
+                        reason: 'tracking',
                     });
                     fileInfo.diagnosticsVersion = undefined;
                 }

packages/pyright-internal/src/backgroundAnalysisBase.ts

@@ -781,6 +781,7 @@ function convertAnalysisResults(result: AnalysisResults): AnalysisResults {
             fileUri: Uri.fromJsonObj(f.fileUri),
             version: f.version,
             diagnostics: convertDiagnostics(f.diagnostics),
+            reason: f.reason,
         };
     });
 
@@ -791,7 +792,7 @@ function convertDiagnostics(diagnostics: Diagnostic[]) {
     // Elements are typed as "any" since data crossing the process
     // boundary loses type info.
     return diagnostics.map<Diagnostic>((d: any) => {
-        const diag = new Diagnostic(d.category, d.message, d.range, d.priority);
+        const diag = new Diagnostic(d.category, d.message, d.range, d.priority, d.baselineStatus);
         if (d._actions) {
             for (const action of d._actions) {
                 diag.addAction(action);

packages/pyright-internal/src/baseline.ts

@@ -157,9 +157,9 @@ export const filterOutBaselinedDiagnostics = (
                             convertLevelToCategory(name),
                             diagnostic.message,
                             diagnostic.range,
-                            diagnostic.priority
+                            diagnostic.priority,
+                            'baselined with hint'
                         );
-                        newDiagnostic.baselineStatus = 'baselined with hint';
                         const rule = diagnostic.getRule();
                         if (rule) {
                             newDiagnostic.setRule(rule);

packages/pyright-internal/src/commands/writeBaseline.ts

@@ -41,7 +41,13 @@ export class WriteBaselineCommand implements ServerCommand {
                 // filter out excluded files. ideally they shouldn't be present at all. see
                 // https://github.com/DetachHead/basedpyright/issues/31
                 const filteredFiles = Object.entries(this._ls.documentsWithDiagnostics)
-                    .filter(([filePath]) => matchFileSpecs(configOptions, Uri.file(filePath, this._ls.serviceProvider)))
+                    .filter(
+                        ([filePath, fileDiagnostics]) =>
+                            // filter out files whose diagnostics were cleared due to the file being closed
+                            fileDiagnostics.reason === 'analysis' &&
+                            // filter out files that aren't included in the project
+                            matchFileSpecs(configOptions, Uri.file(filePath, this._ls.serviceProvider))
+                    )
                     .map(([_, diagnostics]) => diagnostics);
                 const newBaseline = writeDiagnosticsToBaselineFile(
                     workspace.service.fs,

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

@@ -124,7 +124,6 @@ export type BaselineStatus = 'baselined' | 'baselined with hint';
 
 // Represents a single error or warning.
 export class Diagnostic {
-    baselineStatus?: BaselineStatus;
     private _actions: DiagnosticAction[] | undefined;
     private _rule: string | undefined;
     private _relatedInfo: DiagnosticRelatedInfo[] = [];
@@ -133,7 +132,8 @@ export class Diagnostic {
         readonly category: DiagnosticCategory,
         readonly message: string,
         readonly range: Range,
-        readonly priority: TaskListPriority = TaskListPriority.Normal
+        readonly priority: TaskListPriority = TaskListPriority.Normal,
+        public baselineStatus: BaselineStatus | undefined = undefined
     ) {}
 
     toJsonObj() {
@@ -142,14 +142,15 @@ export class Diagnostic {
             message: this.message,
             range: this.range,
             priority: this.priority,
+            baselineStatus: this.baselineStatus,
             actions: this._actions,
             rule: this._rule,
             relatedInfo: this._relatedInfo.map((info) => DiagnosticRelatedInfo.toJsonObj(info)),
         };
     }
 
     static fromJsonObj(obj: any) {
-        const diag = new Diagnostic(obj.category, obj.message, obj.range, obj.priority);
+        const diag = new Diagnostic(obj.category, obj.message, obj.range, obj.priority, obj.baselineStatus);
         diag._actions = obj.actions;
         diag._rule = obj.rule;
         diag._relatedInfo = obj.relatedInfo.map((info: any) => DiagnosticRelatedInfo.fromJsonObj(info));

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

@@ -22,6 +22,7 @@ export interface FileDiagnostics {
     fileUri: Uri;
     version: number | undefined;
     diagnostics: Diagnostic[];
+    reason: 'analysis' | 'tracking';
 }
 
 export namespace FileDiagnostics {
@@ -30,6 +31,7 @@ export namespace FileDiagnostics {
             fileUri: fileDiag.fileUri.toJsonObj(),
             version: fileDiag.version,
             diagnostics: fileDiag.diagnostics.map((d) => d.toJsonObj()),
+            reason: fileDiag.reason,
         };
     }
 
@@ -38,6 +40,7 @@ export namespace FileDiagnostics {
             fileUri: Uri.fromJsonObj(fileDiagObj.fileUri),
             version: fileDiagObj.version,
             diagnostics: fileDiagObj.diagnostics.map((d: any) => Diagnostic.fromJsonObj(d)),
+            reason: fileDiagObj.reason,
         };
     }
 }

packages/pyright-internal/src/languageServerBase.ts

@@ -192,7 +192,7 @@ export abstract class LanguageServerBase implements LanguageServerInterface, Dis
     >();
 
     // The URIs for which diagnostics are reported
-    readonly documentsWithDiagnostics: Record<string, FileDiagnostics & { reason: 'analysis' | 'tracking' }> = {};
+    readonly documentsWithDiagnostics: Record<string, FileDiagnostics> = {};
 
     protected readonly dynamicFeatures = new DynamicFeatures();
 
@@ -1276,7 +1276,13 @@ export abstract class LanguageServerBase implements LanguageServerInterface, Dis
         );
 
         // if any baselined diagnostics disappeared, update the baseline for the effected files
-        if (results.reason === 'analysis') {
+        if (
+            results.reason === 'analysis' &&
+            // if there are still any files requirign analysis, don't update the baseline file as it could
+            // incorrectly delete diagnostics that are still present
+            !results.requiringAnalysisCount.files &&
+            !results.requiringAnalysisCount.cells
+        ) {
             const filesRequiringBaselineUpdate = new Map<Workspace, FileDiagnostics[]>();
             for (const [fileUri, savedFileInfo] of this.savedFilesForBaselineUpdate.entries()) {
                 // can't use result.diagnostics because we need the diagnostics from the previous analysis since
@@ -1285,24 +1291,21 @@ export abstract class LanguageServerBase implements LanguageServerInterface, Dis
                 if (!fileDiagnostics || fileDiagnostics.reason !== 'analysis') {
                     continue;
                 }
-                const sourceFile = savedFileInfo.workspace.service.getSourceFile(fileUri);
-                if (sourceFile && !sourceFile.isCheckingRequired()) {
-                    const baselineInfo = getBaselinedErrorsForFile(this.fs, savedFileInfo.workspace.rootUri!, fileUri);
-                    if (
-                        // no baseline file exists or no baselined errors exist for this file
-                        !baselineInfo.length ||
-                        // there are diagnostics that haven't been baselined, so we don't want to write them
-                        // because the user will have to either fix the diagnostics or explicitly write them to the
-                        // baseline themselves
-                        fileDiagnostics.diagnostics.some((diagnostic) => !diagnostic.baselineStatus)
-                    ) {
-                        continue;
-                    }
-                    if (!filesRequiringBaselineUpdate.has(savedFileInfo.workspace)) {
-                        filesRequiringBaselineUpdate.set(savedFileInfo.workspace, []);
-                    }
-                    filesRequiringBaselineUpdate.get(savedFileInfo.workspace)!.push(fileDiagnostics);
+                const baselineInfo = getBaselinedErrorsForFile(this.fs, savedFileInfo.workspace.rootUri!, fileUri);
+                if (
+                    // no baseline file exists or no baselined errors exist for this file
+                    !baselineInfo.length ||
+                    // there are diagnostics that haven't been baselined, so we don't want to write them
+                    // because the user will have to either fix the diagnostics or explicitly write them to the
+                    // baseline themselves
+                    fileDiagnostics.diagnostics.some((diagnostic) => !diagnostic.baselineStatus)
+                ) {
+                    continue;
+                }
+                if (!filesRequiringBaselineUpdate.has(savedFileInfo.workspace)) {
+                    filesRequiringBaselineUpdate.set(savedFileInfo.workspace, []);
                 }
+                filesRequiringBaselineUpdate.get(savedFileInfo.workspace)!.push(fileDiagnostics);
             }
             for (const [workspace, files] of filesRequiringBaselineUpdate.entries()) {
                 writeDiagnosticsToBaselineFile(this.fs, workspace.rootUri!, files, true);
@@ -1365,7 +1368,7 @@ export abstract class LanguageServerBase implements LanguageServerInterface, Dis
                 }
                 this.sendDiagnostics(this.fs, {
                     fileUri: fileWithDiagnostics.fileUri,
-                    diagnostics: [],
+                    diagnostics: fileWithDiagnostics.diagnostics,
                     version: undefined,
                     reason: 'tracking',
                 });