rework lsp notebook support to support baseline

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

@@ -348,26 +348,25 @@ export class Program {
                     importName,
                     isThirdPartyImport,
                     isInPyTypedPackage,
-                    this.baselineHandler,
                     this._editModeTracker,
+                    this.baselineHandler,
+                    () => sourceFileInfo.cellIndex(),
                     this._console,
                     this._logTracker,
                     IPythonMode.CellDocs
                 );
-                sourceFile.setCellIndex(index);
-                sourceFileInfos.push(
-                    new SourceFileInfo(
-                        sourceFile,
-                        isTypeshedFile,
-                        isThirdPartyImport,
-                        isInPyTypedPackage,
-                        this._editModeTracker,
-                        {
-                            isTracked: true,
-                            chainedSourceFile: sourceFileInfos[index - 1],
-                        }
-                    )
+                const sourceFileInfo = new SourceFileInfo(
+                    sourceFile,
+                    isTypeshedFile,
+                    isThirdPartyImport,
+                    isInPyTypedPackage,
+                    this._editModeTracker,
+                    {
+                        isTracked: true,
+                        chainedSourceFile: sourceFileInfos[index - 1],
+                    }
                 );
+                sourceFileInfos.push(sourceFileInfo);
             });
         } else {
             const importName = this._getImportNameForNewSourceFile(fileUri);
@@ -380,8 +379,9 @@ export class Program {
                 importName,
                 isThirdPartyImport,
                 isInPyTypedPackage,
-                this.baselineHandler,
                 this._editModeTracker,
+                this.baselineHandler,
+                () => undefined,
                 this._console,
                 this._logTracker
             );
@@ -412,8 +412,9 @@ export class Program {
                 moduleImportInfo.moduleName,
                 /* isThirdPartyImport */ false,
                 moduleImportInfo.isThirdPartyPyTypedPresent,
-                this.baselineHandler,
                 this._editModeTracker,
+                this.baselineHandler,
+                () => sourceFileInfo?.cellIndex(),
                 this._console,
                 this._logTracker,
                 options?.ipythonMode ?? IPythonMode.None
@@ -1536,8 +1537,9 @@ export class Program {
                         moduleImportInfo.moduleName,
                         importInfo.isThirdPartyImport,
                         importInfo.isPyTypedPresent,
-                        this.baselineHandler,
                         this._editModeTracker,
+                        this.baselineHandler,
+                        () => importedFileInfo?.cellIndex(),
                         this._console,
                         this._logTracker
                     );
@@ -1679,8 +1681,9 @@ export class Program {
             moduleImportInfo.moduleName,
             /* isThirdPartyImport */ false,
             /* isInPyTypedPackage */ false,
-            this.baselineHandler,
             this._editModeTracker,
+            this.baselineHandler,
+            () => sourceFileInfo.cellIndex(),
             this._console,
             this._logTracker
         );

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

@@ -19,8 +19,9 @@ export interface ISourceFileFactory {
         moduleName: string,
         isThirdPartyImport: boolean,
         isThirdPartyPyTypedPresent: boolean,
-        baselineHandler: BaselineHandler,
         editMode: SourceFileEditMode,
+        baselineHandler: BaselineHandler,
+        getCellIndex: () => number | undefined,
         console?: ConsoleInterface,
         logTracker?: LogTracker,
         ipythonMode?: IPythonMode

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

@@ -203,11 +203,6 @@ class WriteableData {
 
     parserOutput: ParserOutput | undefined;
 
-    /**
-     * this is only writable in the language server because when the user moves cells around, the index changes
-     */
-    cellIndex: number | undefined;
-
     constructor() {
         // Empty
     }
@@ -317,6 +312,9 @@ export class SourceFile {
         isThirdPartyPyTypedPresent: boolean,
         editMode: SourceFileEditMode,
         private _baselineHandler: BaselineHandler,
+        // this is kinda weird but it's necessary because the chained file is stored on SourceFileInfo and
+        // accessing it from here would cause a circular dependency
+        private _getCellIndex: () => number | undefined,
         console?: ConsoleInterface,
         logTracker?: LogTracker,
         ipythonMode?: IPythonMode
@@ -539,10 +537,6 @@ export class SourceFile {
         return this._writableData.clientDocumentContents;
     }
 
-    // TODO: this sucks. ideally SourceFile would just have access to its chained source files, but there's a bunch of machinery
-    // around writable data that i'm too scared to touch
-    setCellIndex = (value: number) => (this._writableData.cellIndex = value);
-
     /**
      * gets the content of the source file. if it's a notebook, the content of this source file's {@link _ipythonCellIndex} is returned
      */
@@ -556,7 +550,7 @@ export class SourceFile {
             // Otherwise, get content from file system.
             return getFileContent(this.fileSystem, this._uri, this._console);
         }
-        const cellIndex = this._writableData.cellIndex;
+        const cellIndex = this._getCellIndex();
         if (cellIndex === undefined) {
             throw new Error(`something went wrong, failed to get cell index for ${this._uri}`);
         }
@@ -1314,11 +1308,7 @@ export class SourceFile {
         // Now add in the "unnecessary type ignore" diagnostics.
         diagList = diagList.concat(unnecessaryTypeIgnoreDiags);
 
-        diagList = this._baselineHandler.sortDiagnosticsAndMatchBaseline(
-            this._uri,
-            this._writableData.cellIndex,
-            diagList
-        );
+        diagList = this._baselineHandler.sortDiagnosticsAndMatchBaseline(this._uri, this._getCellIndex(), diagList);
 
         // If we're not returning any diagnostics, filter out all of
         // the errors and warnings, leaving only the unreachable code

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

@@ -123,7 +123,6 @@ export class SourceFileInfo {
             result++;
             chainedFile = chainedFile.chainedSourceFile;
         }
-        this.sourceFile.setCellIndex(result);
         return result;
     };
 

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

@@ -32,7 +32,11 @@ export class QuickActionCommand implements ServerCommand {
                 return performQuickAction(p, docUri, params.command, otherArgs, token);
             }, token);
 
-            return convertToWorkspaceEdit(workspace.service.fs, convertToFileTextEdits(docUri, editActions ?? []));
+            return convertToWorkspaceEdit(
+                this._ls,
+                workspace.service.fs,
+                convertToFileTextEdits(docUri, editActions ?? [])
+            );
         }
     }
 }

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

@@ -15,7 +15,7 @@ import { DiagnosticBooleanOverridesMap, DiagnosticSeverityOverridesMap } from '.
 import { SignatureDisplayType } from './configOptions';
 import { ConsoleInterface, LogLevel } from './console';
 import { TaskListToken } from './diagnostic';
-import { FileSystem } from './fileSystem';
+import { FileSystem, ReadOnlyFileSystem } from './fileSystem';
 import { FileWatcherHandler } from './fileWatcher';
 import { ServiceProvider } from './serviceProvider';
 import { Uri } from './uri/uri';
@@ -138,6 +138,7 @@ export interface LanguageServerBaseInterface {
 
 export interface LanguageServerInterface extends LanguageServerBaseInterface {
     getWorkspaceForFile(fileUri: Uri, pythonPath?: Uri): Promise<Workspace>;
+    convertUriToLspUriString: (fs: ReadOnlyFileSystem, uri: Uri) => string;
     readonly documentsWithDiagnostics: Record<string, FileDiagnostics>;
 }
 

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

@@ -105,8 +105,9 @@ const DefaultSourceFileFactory: ISourceFileFactory = {
         moduleName: string,
         isThirdPartyImport: boolean,
         isThirdPartyPyTypedPresent: boolean,
-        baselineHandler: BaselineHandler,
         editMode: SourceFileEditMode,
+        baselineHandler: BaselineHandler,
+        getCellIndex: () => number | undefined,
         console?: ConsoleInterface,
         logTracker?: LogTracker,
         ipythonMode?: IPythonMode
@@ -119,6 +120,7 @@ const DefaultSourceFileFactory: ISourceFileFactory = {
             isThirdPartyPyTypedPresent,
             editMode,
             baselineHandler,
+            getCellIndex,
             console,
             logTracker,
             ipythonMode

packages/pyright-internal/src/common/uri/uriUtils.ts

@@ -276,7 +276,7 @@ export function getWildcardRoot(root: Uri, fileSpec: string): Uri {
 }
 
 export function hasPythonExtension(uri: Uri) {
-    return uri.hasExtension('.py') || uri.hasExtension('.pyi');
+    return uri.hasExtension('.py') || uri.hasExtension('.pyi') || uri.hasExtension('.ipynb');
 }
 
 export function getFileSpec(root: Uri, fileSpec: string): FileSpec {
@@ -386,11 +386,6 @@ export function getRootUri(csdOrSp: CaseSensitivityDetector | ServiceProvider):
     return undefined;
 }
 
-export function convertUriToLspUriString(fs: ReadOnlyFileSystem, uri: Uri): string {
-    // Convert to a URI string that the LSP client understands (mapped files are only local to the server).
-    return fs.getOriginalUri(uri).toString();
-}
-
 export namespace UriEx {
     export function file(path: string): Uri;
     export function file(path: string, isCaseSensitive: boolean, checkRelative?: boolean): Uri;

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

@@ -28,7 +28,7 @@ import { convertRangeToTextRange, convertTextRangeToRange } from './positionUtil
 import { TextRange } from './textRange';
 import { TextRangeCollection } from './textRangeCollection';
 import { Uri } from './uri/uri';
-import { convertUriToLspUriString } from './uri/uriUtils';
+import { LanguageServerInterface } from './languageServerInterface';
 
 export function convertToTextEdits(editActions: TextEditAction[]): TextEdit[] {
     return editActions.map((editAction) => ({
@@ -47,9 +47,18 @@ export function convertToFileTextEdits(fileUri: Uri, editActions: TextEditAction
     return editActions.map((a) => ({ fileUri, ...a }));
 }
 
-export function convertToWorkspaceEdit(fs: ReadOnlyFileSystem, edits: FileEditAction[]): WorkspaceEdit;
-export function convertToWorkspaceEdit(fs: ReadOnlyFileSystem, edits: FileEditActions): WorkspaceEdit;
 export function convertToWorkspaceEdit(
+    ls: LanguageServerInterface,
+    fs: ReadOnlyFileSystem,
+    edits: FileEditAction[]
+): WorkspaceEdit;
+export function convertToWorkspaceEdit(
+    ls: LanguageServerInterface,
+    fs: ReadOnlyFileSystem,
+    edits: FileEditActions
+): WorkspaceEdit;
+export function convertToWorkspaceEdit(
+    ls: LanguageServerInterface,
     fs: ReadOnlyFileSystem,
     edits: FileEditActions,
     changeAnnotations: {
@@ -58,6 +67,7 @@ export function convertToWorkspaceEdit(
     defaultAnnotationId: string
 ): WorkspaceEdit;
 export function convertToWorkspaceEdit(
+    ls: LanguageServerInterface,
     fs: ReadOnlyFileSystem,
     edits: FileEditActions | FileEditAction[],
     changeAnnotations?: {
@@ -66,15 +76,20 @@ export function convertToWorkspaceEdit(
     defaultAnnotationId = 'default'
 ): WorkspaceEdit {
     if (isArray(edits)) {
-        return _convertToWorkspaceEditWithChanges(fs, edits);
+        return _convertToWorkspaceEditWithChanges(ls, fs, edits);
     }
 
-    return _convertToWorkspaceEditWithDocumentChanges(fs, edits, changeAnnotations, defaultAnnotationId);
+    return _convertToWorkspaceEditWithDocumentChanges(ls, fs, edits, changeAnnotations, defaultAnnotationId);
 }
 
-export function appendToWorkspaceEdit(fs: ReadOnlyFileSystem, edits: FileEditAction[], workspaceEdit: WorkspaceEdit) {
+export function appendToWorkspaceEdit(
+    ls: LanguageServerInterface,
+    fs: ReadOnlyFileSystem,
+    edits: FileEditAction[],
+    workspaceEdit: WorkspaceEdit
+) {
     edits.forEach((edit) => {
-        const uri = convertUriToLspUriString(fs, edit.fileUri);
+        const uri = ls.convertUriToLspUriString(fs, edit.fileUri);
         workspaceEdit.changes![uri] = workspaceEdit.changes![uri] || [];
         workspaceEdit.changes![uri].push({ range: edit.range, newText: edit.replacementText });
     });
@@ -167,6 +182,7 @@ export function applyDocumentChanges(program: EditableProgram, fileInfo: SourceF
 }
 
 export function generateWorkspaceEdit(
+    ls: LanguageServerInterface,
     fs: ReadOnlyFileSystem,
     originalService: AnalyzerService,
     clonedService: AnalyzerService,
@@ -191,7 +207,7 @@ export function generateWorkspaceEdit(
             continue;
         }
 
-        edits.changes![convertUriToLspUriString(fs, uri)] = [
+        edits.changes![ls.convertUriToLspUriString(fs, uri)] = [
             {
                 range: convertTextRangeToRange(parseResults.parserOutput.parseTree, parseResults.tokenizerOutput.lines),
                 newText: final.getFileContent() ?? '',
@@ -202,16 +218,21 @@ export function generateWorkspaceEdit(
     return edits;
 }
 
-function _convertToWorkspaceEditWithChanges(fs: ReadOnlyFileSystem, edits: FileEditAction[]) {
+function _convertToWorkspaceEditWithChanges(
+    ls: LanguageServerInterface,
+    fs: ReadOnlyFileSystem,
+    edits: FileEditAction[]
+) {
     const workspaceEdit: WorkspaceEdit = {
         changes: {},
     };
 
-    appendToWorkspaceEdit(fs, edits, workspaceEdit);
+    appendToWorkspaceEdit(ls, fs, edits, workspaceEdit);
     return workspaceEdit;
 }
 
 function _convertToWorkspaceEditWithDocumentChanges(
+    ls: LanguageServerInterface,
     fs: ReadOnlyFileSystem,
     editActions: FileEditActions,
     changeAnnotations?: {
@@ -231,7 +252,7 @@ function _convertToWorkspaceEditWithDocumentChanges(
             case 'create':
                 workspaceEdit.documentChanges!.push(
                     CreateFile.create(
-                        convertUriToLspUriString(fs, operation.fileUri),
+                        ls.convertUriToLspUriString(fs, operation.fileUri),
                         /* options */ undefined,
                         defaultAnnotationId
                     )
@@ -246,7 +267,7 @@ function _convertToWorkspaceEditWithDocumentChanges(
     }
 
     // Text edit's file path must refer to original file paths unless it is a new file just created.
-    const mapPerFile = createMapFromItems(editActions.edits, (e) => convertUriToLspUriString(fs, e.fileUri));
+    const mapPerFile = createMapFromItems(editActions.edits, (e) => ls.convertUriToLspUriString(fs, e.fileUri));
     for (const [uri, value] of mapPerFile) {
         workspaceEdit.documentChanges!.push(
             TextDocumentEdit.create(
@@ -269,8 +290,8 @@ function _convertToWorkspaceEditWithDocumentChanges(
             case 'rename':
                 workspaceEdit.documentChanges!.push(
                     RenameFile.create(
-                        convertUriToLspUriString(fs, operation.oldFileUri),
-                        convertUriToLspUriString(fs, operation.newFileUri),
+                        ls.convertUriToLspUriString(fs, operation.oldFileUri),
+                        ls.convertUriToLspUriString(fs, operation.newFileUri),
                         /* options */ undefined,
                         defaultAnnotationId
                     )
@@ -279,7 +300,7 @@ function _convertToWorkspaceEditWithDocumentChanges(
             case 'delete':
                 workspaceEdit.documentChanges!.push(
                     DeleteFile.create(
-                        convertUriToLspUriString(fs, operation.fileUri),
+                        ls.convertUriToLspUriString(fs, operation.fileUri),
                         /* options */ undefined,
                         defaultAnnotationId
                     )