Skip to content

Commit

Permalink
rework lsp notebook support to support baseline
Browse files Browse the repository at this point in the history
  • Loading branch information
@DetachHead
DetachHead committed Feb 6, 2025
1 parent e97d54a commit 2b13e8d
Showing 1 changed file with 110 additions and 78 deletions.
188 changes: 110 additions & 78 deletions packages/pyright-internal/src/languageServerBase.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,6 @@ import {
InlayHintParams,
SemanticTokens,
SemanticTokensParams,
TextDocumentItem,
WillSaveTextDocumentParams,
} from 'vscode-languageserver-protocol';
import { ResultProgressReporter } from 'vscode-languageserver';
Expand Down Expand Up @@ -151,6 +150,7 @@ import { RenameUsageFinder } from './analyzer/renameUsageFinder';
import { BaselineHandler } from './baseline';
import { assert } from './common/debug';
import { AutoImporter, buildModuleSymbolsMap } from './languageService/autoImporter';
import { zip } from 'lodash';

export abstract class LanguageServerBase<T extends FileWatcherHandler = FileWatcherHandler>
implements LanguageServerInterface, Disposable
Expand Down Expand Up @@ -196,6 +196,7 @@ export abstract class LanguageServerBase<T extends FileWatcherHandler = FileWatc

protected readonly workspaceFactory: WorkspaceFactory;
protected readonly openFileMap = new Map<string, TextDocument>();
private readonly _openCells: TextDocument[] = [];
protected readonly fs: FileSystem;
protected readonly caseSensitiveDetector: CaseSensitivityDetector;

Expand Down Expand Up @@ -267,6 +268,7 @@ export abstract class LanguageServerBase<T extends FileWatcherHandler = FileWatc
dispose() {
this.workspaceFactory.clear();
this.openFileMap.clear();
this._openCells.length = 0;
this.dynamicFeatures.unregister();
this._workspaceFoldersChangedDisposable?.dispose();
}
Expand Down Expand Up @@ -1126,20 +1128,8 @@ export abstract class LanguageServerBase<T extends FileWatcherHandler = FileWatc
}, token);
}

// these overloads ban specifying chainedFile unless iPythonMode is CellDocs
protected async onDidOpenTextDocument(
params: DidOpenTextDocumentParams,
iPythonMode: IPythonMode.CellDocs,
chainedFile?: TextDocumentItem
): Promise<void>;
protected async onDidOpenTextDocument(params: DidOpenTextDocumentParams, iPythonMode?: IPythonMode): Promise<void>;
protected async onDidOpenTextDocument(
params: DidOpenTextDocumentParams,
iPythonMode = IPythonMode.None,
chainedFile?: TextDocumentItem
) {
protected async onDidOpenTextDocument(params: DidOpenTextDocumentParams) {
const uri = this.convertLspUriStringToUri(params.textDocument.uri);

let doc = this.openFileMap.get(uri.key);
if (doc) {
// We shouldn't get an open text document request for an already-opened doc.
Expand All @@ -1154,31 +1144,38 @@ export abstract class LanguageServerBase<T extends FileWatcherHandler = FileWatc
);
}
this.openFileMap.set(uri.key, doc);
const chainedFileUri = chainedFile ? Uri.parse(chainedFile.uri, this.serviceProvider) : undefined;

// Send this open to all the workspaces that might contain this file.
const workspaces = await this.getContainingWorkspacesForFile(uri);
workspaces.forEach((w) => {
w.service.setFileOpened(
uri,
params.textDocument.version,
params.textDocument.text,
iPythonMode,
chainedFileUri
);
w.service.setFileOpened(uri, params.textDocument.version, params.textDocument.text);
});
}

protected onDidOpenNotebookDocument = async (params: DidOpenNotebookDocumentParams) => {
await Promise.all(
params.cellTextDocuments.map((textDocument, index) =>
// the previous cell is the chained document
this.onDidOpenTextDocument({ textDocument }, IPythonMode.CellDocs, params.cellTextDocuments[index - 1])
)
params.cellTextDocuments.map(async (textDocument, index) => {
const uri = this.convertLspUriStringToUri(textDocument.uri, index);
const doc = TextDocument.create(textDocument.uri, 'python', textDocument.version, textDocument.text);
this._openCells.push(doc);
// Send this open to all the workspaces that might contain this file.
const workspaces = await this.getContainingWorkspacesForFile(uri);
workspaces.forEach((w) => {
w.service.setFileOpened(
uri,
textDocument.version,
textDocument.text,
IPythonMode.CellDocs,
// if index > 0 then it has a chained file, which should always just be the previous index because we update them
// on every change
index ? this.convertLspUriStringToUri(textDocument.uri, index - 1) : undefined
);
});
})
);
};

protected async onDidChangeTextDocument(params: DidChangeTextDocumentParams, ipythonMode = IPythonMode.None) {
protected async onDidChangeTextDocument(params: DidChangeTextDocumentParams) {
this.recordUserInteractionTime();

const uri = this.convertLspUriStringToUri(params.textDocument.uri);
Expand All @@ -1195,68 +1192,98 @@ export abstract class LanguageServerBase<T extends FileWatcherHandler = FileWatc
// Send this change to all the workspaces that might contain this file.
const workspaces = await this.getContainingWorkspacesForFile(uri);
workspaces.forEach((w) => {
w.service.updateOpenFileContents(uri, params.textDocument.version, newContents, ipythonMode);
w.service.updateOpenFileContents(uri, params.textDocument.version, newContents);
});
}

protected onDidChangeNotebookDocument = async (params: DidChangeNotebookDocumentParams) => {
this.recordUserInteractionTime();
const changeStructure = params.change.cells?.structure;
// open any new documents first before we action the cell changes, because that's where we attach
// chained documents and if we try to do that before the document is opened it won't work
if (changeStructure?.didOpen) {
await Promise.all(
changeStructure.didOpen.map((textDocument) =>
this.onDidOpenTextDocument({ textDocument }, IPythonMode.CellDocs)
)
);
}
// the rest of these methods can be executed in any order so we collect all their promises and await
// them at the end
const promises: Promise<void>[] = [];
if (changeStructure?.didClose) {
promises.push(
...changeStructure.didClose.map((textDocument) => this.onDidCloseTextDocument({ textDocument }))
if (changeStructure) {
const previousCells = [...this._openCells];
this._openCells.splice(
changeStructure.array.start,
changeStructure.array.deleteCount,
...(changeStructure.array.cells?.map((changedTextDocumentItem) => {
// if there isn't a cell at this index already, we need to open it
const newDocumentItem = changeStructure.didOpen?.find(
(newDocument) => newDocument.uri === changedTextDocumentItem.document
);
if (newDocumentItem) {
return TextDocument.create(
newDocumentItem.uri,
'python',
newDocumentItem.version,
newDocumentItem.text
);
} else {
const result = this._openCells.find(
(openCell) => openCell.uri === changedTextDocumentItem.document
);
if (!result) {
throw new Error(`failed to find existing cell ${changedTextDocumentItem.document}`);
}
return result;
}
}) ?? [])
);
}
const cellChanges = changeStructure?.array.cells;
if (cellChanges) {
promises.push(
...cellChanges.map(async (cell, index) => {
const uri = this.convertLspUriStringToUri(cell.document);
const doc = this.openFileMap.get(uri.key);
if (!doc) {
// We shouldn't get a change text request for a closed doc.
this.console.error(`Received change notebook document command for closed cell ${uri}`);
await Promise.all(
zip(previousCells, this._openCells).map(async ([previousCell, newCell], index) => {
if (previousCell?.uri === newCell?.uri) {
return;
}
// Send this change to all the workspaces that might contain this file.
const workspaces = await this.getContainingWorkspacesForFile(uri);
const previousCell = cellChanges[index - 1]?.document;
const previousCellUri = previousCell ? Uri.parse(previousCell, this.serviceProvider) : undefined;
workspaces.forEach((w) => {
w.service.updateChainedUri(uri, previousCellUri);
});
if (previousCell === undefined) {
// a new cell was added and we didn't already have one at this index so we need to open it
if (newCell === undefined) {
// this should never happen
throw new Error('new cell was undefined when new cell was added');
}
const uri = this.convertLspUriStringToUri(newCell.uri, index);
// Send this open to all the workspaces that might contain this file.
const workspaces = await this.getContainingWorkspacesForFile(uri);
workspaces.forEach((w) => {
w.service.setFileOpened(
uri,
newCell.version,
newCell.getText(),
IPythonMode.CellDocs,
// if cellIndex > 0 then it has a chained file, which should always just be the previous index because we update them
// on every change
index ? this.convertLspUriStringToUri(newCell.uri, index - 1) : undefined
);
});
this.openFileMap.set(uri.key, newCell);
} else {
const uri = this.convertLspUriStringToUri(previousCell.uri, index);
const workspaces = await this.getContainingWorkspacesForFile(uri);
if (newCell === undefined) {
// a cell was deleted and there's no longer a cell at this index so we need to close it
// Send this close to all the workspaces that might contain this file.
workspaces.forEach((w) => w.service.setFileClosed(uri));
this.openFileMap.delete(uri.key);
} else {
// this index now has a different cell than it did before (ie. order was changed) so we have to update the already opened cell
// with the new text content
const newContents = newCell.getText();
TextDocument.update(previousCell, [{ text: newContents }], previousCell.version + 1);
// Send this change to all the workspaces that might contain this file.
workspaces.forEach((w) => {
w.service.updateOpenFileContents(
uri,
previousCell.version + 1,
newContents,
IPythonMode.CellDocs
);
});
}
}
})
);
}
if (params.change.cells?.textContent) {
promises.push(
...params.change.cells.textContent.map((textContentChange) =>
this.onDidChangeTextDocument(
{
textDocument: textContentChange.document,
contentChanges: textContentChange.changes,
},
IPythonMode.CellDocs
)
)
);
}
await Promise.all(promises);
};

protected async onDidCloseTextDocument(params: DidCloseTextDocumentParams) {
const uri = this.convertLspUriStringToUri(params.textDocument.uri);
protected async onDidCloseTextDocument(params: DidCloseTextDocumentParams, cellIndex?: number) {
const uri = this.convertLspUriStringToUri(params.textDocument.uri, cellIndex);

// Send this close to all the workspaces that might contain this file.
const workspaces = await this.getContainingWorkspacesForFile(uri);
Expand Down Expand Up @@ -1389,6 +1416,7 @@ export abstract class LanguageServerBase<T extends FileWatcherHandler = FileWatc

// Stop tracking all open files.
this.openFileMap.clear();
this._openCells.length = 0;
this.serviceProvider.dispose();

return Promise.resolve();
Expand Down Expand Up @@ -1589,8 +1617,12 @@ export abstract class LanguageServerBase<T extends FileWatcherHandler = FileWatc
this.connection.sendDiagnostics(await this.convertDiagnostics(fs, fileWithDiagnostics));
}

protected convertLspUriStringToUri(uri: string) {
return Uri.parse(uri, this.serverOptions.serviceProvider);
protected convertLspUriStringToUri(uri: string, cellIndex?: number | undefined) {
const result = Uri.parse(uri, this.serverOptions.serviceProvider);
if (cellIndex === undefined) {
return result;
}
return Uri.parse(result.getPath(), this.serviceProvider).withFragment(cellIndex.toString());
}

protected addDynamicFeature(feature: DynamicFeature) {
Expand Down

0 comments on commit 2b13e8d

Please sign in to comment.