From 0f647a7eb171850e9eaa698d9320d659c2d5843f Mon Sep 17 00:00:00 2001
From: detachhead <detachhead@users.noreply.github.com>
Date: Tue, 4 Feb 2025 19:34:02 +1000
Subject: [PATCH] prevent invalid baseline file from being generated when
 project contains notebooks (see
 https://github.com/DetachHead/basedpyright/issues/1039)

---
 packages/pyright-internal/src/analyzer/sourceFile.ts    | 9 +++------
 packages/pyright-internal/src/baseline.ts               | 8 ++++++++
 packages/pyright-internal/src/commands/writeBaseline.ts | 9 +++++++--
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/packages/pyright-internal/src/analyzer/sourceFile.ts b/packages/pyright-internal/src/analyzer/sourceFile.ts
index aa59c7e6d..12f731715 100644
--- a/packages/pyright-internal/src/analyzer/sourceFile.ts
+++ b/packages/pyright-internal/src/analyzer/sourceFile.ts
@@ -1005,13 +1005,10 @@ export class SourceFile {
     }
 
     /**
-     * if this source file represents an ipython cell and we're in the cli instead of the language server,
-     * we need to create a fake uri to distinguish between them.
+     * gets the uri of the file on disk. ie. if this source file is a notebook cell, its fragment
+     * that represents its cell id and its `vscode-notebook-cell://` scheme is removed
      */
-    getRealUri = () =>
-        // when using the language server, a fake uri is assigned automatically and this method *should* never get called
-        // because we don't read the files from the disk directly, but we check this anyway just in case
-        this._writableData.clientDocumentContents ? this._uri : this._uri.withFragment('');
+    getRealUri = () => this._uri.withFragment('');
 
     test_enableIPythonMode(enable: boolean) {
         this._ipythonMode = enable ? IPythonMode.CellDocs : IPythonMode.None;
diff --git a/packages/pyright-internal/src/baseline.ts b/packages/pyright-internal/src/baseline.ts
index 3bb452705..4d086a165 100644
--- a/packages/pyright-internal/src/baseline.ts
+++ b/packages/pyright-internal/src/baseline.ts
@@ -78,6 +78,10 @@ class BaselineDiff<T extends boolean> {
     };
 }
 
+export const notebooksNotSupportedMessage =
+    'jupyter notebooks were detected in your project which are not yet supported in baseline files.' +
+    ' see https://github.com/DetachHead/basedpyright/issues/1039';
+
 export class BaselineHandler {
     /**
      * project root can change and we need to invalidate the cache when that happens
@@ -141,6 +145,10 @@ export class BaselineHandler {
                 return undefined;
             }
         }
+        if (filesWithDiagnostics.some((FileDiagnostics) => FileDiagnostics.fileUri.hasExtension('.ipynb'))) {
+            this._console.error(notebooksNotSupportedMessage);
+            return;
+        }
         const newBaselineFiles = this._filteredDiagnosticsToBaselineFormat(filesWithDiagnostics).files;
         const previousBaselineFiles = baselineData?.files ?? {};
         // we don't know for sure that basedpyright was run on every file that was included when the previous baseline was
diff --git a/packages/pyright-internal/src/commands/writeBaseline.ts b/packages/pyright-internal/src/commands/writeBaseline.ts
index 50a228bd6..cb39282b4 100644
--- a/packages/pyright-internal/src/commands/writeBaseline.ts
+++ b/packages/pyright-internal/src/commands/writeBaseline.ts
@@ -1,5 +1,5 @@
 import { ServerCommand } from './commandController';
-import { baselineFilePath, BaselineHandler } from '../baseline';
+import { baselineFilePath, BaselineHandler, notebooksNotSupportedMessage } from '../baseline';
 import { LanguageServerInterface } from '../common/languageServerInterface';
 import { matchFileSpecs } from '../common/configOptions';
 import { Uri } from '../common/uri/uri';
@@ -20,10 +20,15 @@ export class WriteBaselineCommand implements ServerCommand {
         let workspace = workspaces.find((workspace) =>
             workspace.rootUri ? workspace.service.fs.existsSync(baselineFilePath(workspace.rootUri)) : false
         );
+        const filesWithDiagnostics = Object.values(this._ls.documentsWithDiagnostics);
+        if (filesWithDiagnostics.some((FileDiagnostics) => FileDiagnostics.fileUri.scheme === 'vscode-notebook-cell')) {
+            this._ls.window.showWarningMessage(notebooksNotSupportedMessage);
+            return;
+        }
         if (!workspace) {
             // if there's no baseline file yet, we do it in an even hackier way, by getting the workspace from
             // any open file that has diagnostics in it.
-            const firstFile = Object.values(this._ls.documentsWithDiagnostics)[0]?.fileUri;
+            const firstFile = filesWithDiagnostics[0]?.fileUri;
             if (firstFile) {
                 workspace = await this._ls.getWorkspaceForFile(firstFile);
             }