Skip to content

Commit 60578ed

Browse files
andreamahandreamah
authored
implement closed-notebook search (microsoft#174287)
1 parent b008409 commit 60578ed

File tree

9 files changed

+532
-129
lines changed

9 files changed

+532
-129
lines changed

‎src/vs/workbench/contrib/search/browser/notebookSearch.ts

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
7+
import { CancellationToken } from 'vs/base/common/cancellation';
8+
import { ResourceSet } from 'vs/base/common/map';
9+
import { ITextQuery, ISearchProgressItem, ISearchComplete } from 'vs/workbench/services/search/common/search';
10+
11+
export const INotebookSearchService = createDecorator<INotebookSearchService>('notebookSearchService');
12+
13+
export interface INotebookSearchService {
14+
15+
readonly _serviceBrand: undefined;
16+
17+
notebookSearch(query: ITextQuery, token: CancellationToken, onProgress?: (result: ISearchProgressItem) => void): Promise<{ completeData: ISearchComplete; scannedFiles: ResourceSet }>;
18+
}

‎src/vs/workbench/contrib/search/browser/notebookSearchContributions.ts

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
import { InstantiationType, registerSingleton } from 'vs/platform/instantiation/common/extensions';
6+
import { ReplacePreviewContentProvider } from 'vs/workbench/contrib/search/browser/replaceService';
7+
import { Registry } from 'vs/platform/registry/common/platform';
8+
import { IWorkbenchContributionsRegistry, Extensions as WorkbenchExtensions } from 'vs/workbench/common/contributions';
9+
import { LifecyclePhase } from 'vs/workbench/services/lifecycle/common/lifecycle';
10+
import { INotebookSearchService } from 'vs/workbench/contrib/search/browser/notebookSearch';
11+
import { NotebookSearchService } from 'vs/workbench/contrib/search/browser/notebookSearchService';
12+
13+
export function registerContributions(): void {
14+
registerSingleton(INotebookSearchService, NotebookSearchService, InstantiationType.Delayed);
15+
Registry.as<IWorkbenchContributionsRegistry>(WorkbenchExtensions.Workbench).registerWorkbenchContribution(ReplacePreviewContentProvider, LifecyclePhase.Starting);
16+
}

‎src/vs/workbench/contrib/search/browser/notebookSearchService.ts

Lines changed: 310 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,310 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
import { streamToBuffer } from 'vs/base/common/buffer';
6+
import { CancellationToken } from 'vs/base/common/cancellation';
7+
import { IRelativePattern } from 'vs/base/common/glob';
8+
import { ResourceSet, ResourceMap } from 'vs/base/common/map';
9+
import { URI } from 'vs/base/common/uri';
10+
import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
11+
import { IFileService } from 'vs/platform/files/common/files';
12+
import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
13+
import { ILogService } from 'vs/platform/log/common/log';
14+
import { IUriIdentityService } from 'vs/platform/uriIdentity/common/uriIdentity';
15+
import { NotebookEditorWidget } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
16+
import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/services/notebookEditorService';
17+
import { NotebookCellTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookCellTextModel';
18+
import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
19+
import { INotebookExclusiveDocumentFilter, NotebookData } from 'vs/workbench/contrib/notebook/common/notebookCommon';
20+
import { INotebookSerializer, INotebookService, SimpleNotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookService';
21+
import { INotebookSearchService } from 'vs/workbench/contrib/search/browser/notebookSearch';
22+
import { IFileMatchWithCells, ICellMatch, CellSearchModel, contentMatchesToTextSearchMatches, webviewMatchesToTextSearchMatches } from 'vs/workbench/contrib/search/browser/searchNotebookHelpers';
23+
import { IEditorResolverService, priorityToRank } from 'vs/workbench/services/editor/common/editorResolverService';
24+
import { ITextQuery, IFileQuery, QueryType, ISearchProgressItem, ISearchComplete, ISearchConfigurationProperties, ISearchService } from 'vs/workbench/services/search/common/search';
25+
import * as arrays from 'vs/base/common/arrays';
26+
import { isNumber } from 'vs/base/common/types';
27+
28+
interface INotebookDataEditInfo {
29+
notebookData: NotebookData;
30+
mTime: number;
31+
}
32+
33+
interface INotebookSearchMatchResults {
34+
results: ResourceMap<IFileMatchWithCells | null>;
35+
limitHit: boolean;
36+
}
37+
38+
class NotebookDataCache {
39+
private _entries: ResourceMap<INotebookDataEditInfo>;
40+
// private _serializer: INotebookSerializer | undefined;
41+
42+
constructor(
43+
@IUriIdentityService private readonly uriIdentityService: IUriIdentityService,
44+
@IFileService private readonly fileService: IFileService,
45+
@INotebookService private readonly notebookService: INotebookService,
46+
@IEditorResolverService private readonly editorResolverService: IEditorResolverService,
47+
) {
48+
this._entries = new ResourceMap<INotebookDataEditInfo>(uri => this.uriIdentityService.extUri.getComparisonKey(uri));
49+
}
50+
51+
private async getSerializer(notebookUri: URI): Promise<INotebookSerializer | undefined> {
52+
const registeredEditorInfo = this.editorResolverService.getEditors(notebookUri);
53+
const priorityEditorInfo = registeredEditorInfo.reduce((acc, val) =>
54+
priorityToRank(acc.priority) > priorityToRank(val.priority) ? acc : val
55+
);
56+
const info = await this.notebookService.withNotebookDataProvider(priorityEditorInfo.id);
57+
if (!(info instanceof SimpleNotebookProviderInfo)) {
58+
return undefined;
59+
}
60+
return info.serializer;
61+
}
62+
63+
async getNotebookData(notebookUri: URI): Promise<NotebookData> {
64+
const mTime = (await this.fileService.stat(notebookUri)).mtime;
65+
66+
const entry = this._entries.get(notebookUri);
67+
68+
if (entry && entry.mTime === mTime) {
69+
return entry.notebookData;
70+
} else {
71+
72+
let _data: NotebookData = {
73+
metadata: {},
74+
cells: []
75+
};
76+
77+
const content = await this.fileService.readFileStream(notebookUri);
78+
const bytes = await streamToBuffer(content.value);
79+
const serializer = await this.getSerializer(notebookUri);
80+
if (!serializer) {
81+
//unsupported
82+
throw new Error(`serializer not initialized`);
83+
}
84+
_data = await serializer.dataToNotebook(bytes);
85+
this._entries.set(notebookUri, { notebookData: _data, mTime });
86+
return _data;
87+
}
88+
}
89+
90+
}
91+
92+
export class NotebookSearchService implements INotebookSearchService {
93+
declare readonly _serviceBrand: undefined;
94+
private _notebookDataCache: NotebookDataCache;
95+
constructor(
96+
@IInstantiationService private readonly instantiationService: IInstantiationService,
97+
@IUriIdentityService private readonly uriIdentityService: IUriIdentityService,
98+
@INotebookEditorService private readonly notebookEditorService: INotebookEditorService,
99+
@ILogService private readonly logService: ILogService,
100+
@INotebookService private readonly notebookService: INotebookService,
101+
@ISearchService private readonly searchService: ISearchService,
102+
@IConfigurationService private readonly configurationService: IConfigurationService,
103+
104+
) {
105+
this._notebookDataCache = this.instantiationService.createInstance(NotebookDataCache);
106+
}
107+
108+
private async runFileQueries(includes: (string)[], token: CancellationToken, textQuery: ITextQuery): Promise<URI[]> {
109+
const promises = includes.map(include => {
110+
const query: IFileQuery = {
111+
type: QueryType.File,
112+
filePattern: include,
113+
folderQueries: textQuery.folderQueries,
114+
maxResults: textQuery.maxResults,
115+
};
116+
return this.searchService.fileSearch(
117+
query,
118+
token
119+
);
120+
});
121+
const result = (await Promise.all(promises)).map(sc => sc.results.map(fm => fm.resource)).flat();
122+
const uris = new ResourceSet(result, uri => this.uriIdentityService.extUri.getComparisonKey(uri));
123+
return Array.from(uris.keys());
124+
}
125+
126+
async notebookSearch(query: ITextQuery, token: CancellationToken, onProgress?: (result: ISearchProgressItem) => void): Promise<{ completeData: ISearchComplete; scannedFiles: ResourceSet }> {
127+
128+
if (query.type !== QueryType.Text) {
129+
return {
130+
completeData: {
131+
messages: [],
132+
limitHit: false,
133+
results: [],
134+
},
135+
scannedFiles: new ResourceSet()
136+
};
137+
}
138+
const searchStart = Date.now();
139+
140+
const localNotebookWidgets = this.getLocalNotebookWidgets();
141+
const localNotebookFiles = localNotebookWidgets.map(widget => widget.viewModel!.uri);
142+
const localResultPromise = this.getLocalNotebookResults(query, token, localNotebookWidgets);
143+
const searchLocalEnd = Date.now();
144+
145+
const experimentalNotebooksEnabled = this.configurationService.getValue<ISearchConfigurationProperties>('search').experimental?.closedNotebookRichContentResults ?? false;
146+
147+
let closedResultsPromise: Promise<INotebookSearchMatchResults | undefined> = Promise.resolve(undefined);
148+
if (experimentalNotebooksEnabled) {
149+
closedResultsPromise = this.getClosedNotebookResults(query, new ResourceSet(localNotebookFiles, uri => this.uriIdentityService.extUri.getComparisonKey(uri)), token);
150+
}
151+
152+
const resolved = (await Promise.all([localResultPromise, closedResultsPromise])).filter((result): result is INotebookSearchMatchResults => !!result);
153+
const resultArray = resolved.map(elem => elem.results);
154+
155+
const results = arrays.coalesce(resultArray.flatMap(map => Array.from(map.values())));
156+
const scannedFiles = new ResourceSet(resultArray.flatMap(map => Array.from(map.keys())), uri => this.uriIdentityService.extUri.getComparisonKey(uri));
157+
if (onProgress) {
158+
results.forEach(onProgress);
159+
}
160+
this.logService.trace(`local notebook search time | ${searchLocalEnd - searchStart}ms`);
161+
return {
162+
completeData: {
163+
messages: [],
164+
limitHit: resolved.reduce((prev, cur) => prev || cur.limitHit, false),
165+
results,
166+
},
167+
scannedFiles
168+
};
169+
}
170+
171+
private async getClosedNotebookResults(textQuery: ITextQuery, scannedFiles: ResourceSet, token: CancellationToken): Promise<INotebookSearchMatchResults> {
172+
const infoProviders = this.notebookService.getContributedNotebookTypes();
173+
const includes = infoProviders.flatMap(
174+
(provider) => {
175+
return provider.selectors.map((selector) => {
176+
const globPattern = (selector as INotebookExclusiveDocumentFilter).include || selector as IRelativePattern | string;
177+
return globPattern.toString();
178+
}
179+
);
180+
}
181+
);
182+
183+
const results = new ResourceMap<IFileMatchWithCells | null>(uri => this.uriIdentityService.extUri.getComparisonKey(uri));
184+
185+
const start = Date.now();
186+
187+
const filesToScan = await this.runFileQueries(includes, token, textQuery);
188+
const deserializedNotebooks = new ResourceMap<NotebookTextModel>();
189+
const textModels = this.notebookService.getNotebookTextModels();
190+
for (const notebook of textModels) {
191+
deserializedNotebooks.set(notebook.uri, notebook);
192+
}
193+
194+
const promises = filesToScan.map(async (uri) => {
195+
const cellMatches: ICellMatch[] = [];
196+
if (scannedFiles.has(uri)) {
197+
return;
198+
}
199+
200+
try {
201+
if (token.isCancellationRequested) {
202+
return;
203+
}
204+
205+
const cells = deserializedNotebooks.get(uri)?.cells ?? (await this._notebookDataCache.getNotebookData(uri)).cells;
206+
207+
if (token.isCancellationRequested) {
208+
return;
209+
}
210+
211+
cells.forEach((cell, index) => {
212+
const target = textQuery.contentPattern.pattern;
213+
const cellModel = cell instanceof NotebookCellTextModel ? new CellSearchModel('', cell.textBuffer, uri, index) : new CellSearchModel(cell.source, undefined, uri, index);
214+
215+
const matches = cellModel.find(target);
216+
if (matches.length > 0) {
217+
const cellMatch: ICellMatch = {
218+
cell: cellModel,
219+
index: index,
220+
contentResults: contentMatchesToTextSearchMatches(matches, cellModel),
221+
webviewResults: []
222+
};
223+
cellMatches.push(cellMatch);
224+
}
225+
});
226+
227+
const fileMatch = cellMatches.length > 0 ? {
228+
resource: uri, cellResults: cellMatches
229+
} : null;
230+
results.set(uri, fileMatch);
231+
return;
232+
233+
} catch (e) {
234+
this.logService.info('error: ' + e);
235+
return;
236+
}
237+
238+
});
239+
240+
await Promise.all(promises);
241+
const end = Date.now();
242+
243+
this.logService.trace(`query: ${textQuery.contentPattern.pattern}`);
244+
this.logService.trace(`closed notebook search time | ${end - start}ms`);
245+
return {
246+
results: results,
247+
limitHit: false
248+
};
249+
}
250+
251+
private async getLocalNotebookResults(query: ITextQuery, token: CancellationToken, widgets: Array<NotebookEditorWidget>): Promise<INotebookSearchMatchResults> {
252+
const localResults = new ResourceMap<IFileMatchWithCells | null>(uri => this.uriIdentityService.extUri.getComparisonKey(uri));
253+
let limitHit = false;
254+
255+
for (const widget of widgets) {
256+
if (!widget.viewModel) {
257+
continue;
258+
}
259+
const askMax = isNumber(query.maxResults) ? query.maxResults + 1 : Number.MAX_SAFE_INTEGER;
260+
let matches = await widget
261+
.find(query.contentPattern.pattern, {
262+
regex: query.contentPattern.isRegExp,
263+
wholeWord: query.contentPattern.isWordMatch,
264+
caseSensitive: query.contentPattern.isCaseSensitive,
265+
includeMarkupInput: query.contentPattern.notebookInfo?.isInNotebookMarkdownInput,
266+
includeMarkupPreview: query.contentPattern.notebookInfo?.isInNotebookMarkdownPreview,
267+
includeCodeInput: query.contentPattern.notebookInfo?.isInNotebookCellInput,
268+
includeOutput: query.contentPattern.notebookInfo?.isInNotebookCellOutput,
269+
}, token, false, true);
270+
271+
272+
if (matches.length) {
273+
if (askMax && matches.length >= askMax) {
274+
limitHit = true;
275+
matches = matches.slice(0, askMax - 1);
276+
}
277+
const cellResults: ICellMatch[] = matches.map(match => {
278+
const contentResults = contentMatchesToTextSearchMatches(match.contentMatches, match.cell);
279+
const webviewResults = webviewMatchesToTextSearchMatches(match.webviewMatches);
280+
return {
281+
cell: match.cell,
282+
index: match.index,
283+
contentResults: contentResults,
284+
webviewResults: webviewResults,
285+
};
286+
});
287+
288+
const fileMatch: IFileMatchWithCells = {
289+
resource: widget.viewModel.uri, cellResults: cellResults
290+
};
291+
localResults.set(widget.viewModel.uri, fileMatch);
292+
} else {
293+
localResults.set(widget.viewModel.uri, null);
294+
}
295+
}
296+
297+
return {
298+
results: localResults,
299+
limitHit
300+
};
301+
}
302+
303+
304+
private getLocalNotebookWidgets(): Array<NotebookEditorWidget> {
305+
const notebookWidgets = this.notebookEditorService.retrieveAllExistingWidgets();
306+
return notebookWidgets
307+
.map(widget => widget.value)
308+
.filter((val): val is NotebookEditorWidget => !!val && !!(val.viewModel));
309+
}
310+
}

‎src/vs/workbench/contrib/search/browser/search.contribution.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import { Extensions as ViewExtensions, IViewContainersRegistry, IViewDescriptor,
2121
import { GotoSymbolQuickAccessProvider } from 'vs/workbench/contrib/codeEditor/browser/quickaccess/gotoSymbolQuickAccess';
2222
import { AnythingQuickAccessProvider } from 'vs/workbench/contrib/search/browser/anythingQuickAccess';
2323
import { registerContributions as replaceContributions } from 'vs/workbench/contrib/search/browser/replaceContributions';
24+
import { registerContributions as notebookSearchContributions } from 'vs/workbench/contrib/search/browser/notebookSearchContributions';
2425
import { searchViewIcon } from 'vs/workbench/contrib/search/browser/searchIcons';
2526
import { SearchView } from 'vs/workbench/contrib/search/browser/searchView';
2627
import { registerContributions as searchWidgetContributions } from 'vs/workbench/contrib/search/browser/searchWidget';
@@ -46,6 +47,7 @@ registerSingleton(ISearchWorkbenchService, SearchWorkbenchService, Instantiation
4647
registerSingleton(ISearchHistoryService, SearchHistoryService, InstantiationType.Delayed);
4748

4849
replaceContributions();
50+
notebookSearchContributions();
4951
searchWidgetContributions();
5052

5153
const SEARCH_MODE_CONFIG = 'search.mode';
@@ -349,6 +351,11 @@ configurationRegistry.registerConfiguration({
349351
nls.localize('scm.defaultViewMode.list', "Shows search results as a list.")
350352
],
351353
'description': nls.localize('search.defaultViewMode', "Controls the default search result view mode.")
354+
},
355+
'search.experimental.closedNotebookRichContentResults': {
356+
type: 'boolean',
357+
description: nls.localize('search.experimental.closedNotebookResults', "Show notebook editor rich content results for closed notebooks. Please refresh your search results after changing this setting."),
358+
default: false
352359
}
353360
}
354361
});

0 commit comments

Comments
 (0)