Skip to content

Commit

Permalink
show closed-notebook matches in search view
Browse files Browse the repository at this point in the history
  • Loading branch information
andreamah committed Apr 12, 2023
1 parent 3765abf commit 2abfedc
Show file tree
Hide file tree
Showing 3 changed files with 164 additions and 23 deletions.
26 changes: 16 additions & 10 deletions src/vs/workbench/contrib/search/browser/searchActionsNotebook.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@ import { NotebookData } from 'vs/workbench/contrib/notebook/common/notebookCommo
import { IFileQuery, ISearchService, QueryType } from 'vs/workbench/services/search/common/search';
import { CancellationToken } from 'vs/base/common/cancellation';
import { category } from 'vs/workbench/contrib/search/browser/searchActionsBase';
import { URI } from 'vs/base/common/uri';
import { CellSearchModel } from 'vs/workbench/contrib/search/browser/searchNotebookHelpers';

registerAction2(class NotebookDeserializeTest extends Action2 {

Expand All @@ -41,10 +43,6 @@ registerAction2(class NotebookDeserializeTest extends Action2 {
const currWorkspace = workspacesService.getWorkspace();
const uri = currWorkspace.folders[0].uri;


// const queryBuilder = instantiationService.createInstance(QueryBuilder);

// const query = queryBuilder.text(content, folderResources.map(folder => folder.uri))
const query: IFileQuery = {
type: QueryType.File,
filePattern: '**/*.ipynb',
Expand All @@ -54,16 +52,16 @@ registerAction2(class NotebookDeserializeTest extends Action2 {
query,
CancellationToken.None
);
// glob(dir + '/**/*.ipynb', {}, async (err, files) => {
logService.info('notebook deserialize START');
let processedFiles = 0;
let processedBytes = 0;
let processedCells = 0;
let matchCount = 0;
const start = Date.now();
// const pattern = "text";
const pattern = 'start_index';
let i = 0;
for (const fileMatch of searchComplete.results) {
if (i > 10) {
if (i > Number.MAX_SAFE_INTEGER) {
break;
}
i++;
Expand All @@ -86,9 +84,10 @@ registerAction2(class NotebookDeserializeTest extends Action2 {
}


_data.cells.forEach(cell => {
_data.cells.forEach((cell, index) => {
const input = cell.source;
logService.info(input);
const matches = this.getMatches(input, uri, index, pattern);
matchCount += matches.length;
});

processedFiles += 1;
Expand All @@ -100,7 +99,14 @@ registerAction2(class NotebookDeserializeTest extends Action2 {
}
}
const end = Date.now();
logService.info(`${matchCount} matches found`);
logService.info(`notebook deserialize END | ${end - start}ms | ${((processedBytes / 1024) / 1024).toFixed(2)}MB | Number of Files: ${processedFiles} | Number of Cells: ${processedCells}`);
// });

}

getMatches(source: string, uri: URI, cellIndex: number, target: string) {
const cellModel = new CellSearchModel(source, uri, cellIndex);
return cellModel.find(target);
}
});

98 changes: 88 additions & 10 deletions src/vs/workbench/contrib/search/browser/searchModel.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import * as arrays from 'vs/base/common/arrays';
import { RunOnceScheduler } from 'vs/base/common/async';
import { streamToBuffer } from 'vs/base/common/buffer';
import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
import { compareFileExtensions, compareFileNames, comparePaths } from 'vs/base/common/comparers';
import { memoize } from 'vs/base/common/decorators';
Expand Down Expand Up @@ -36,11 +37,12 @@ import { FindMatchDecorationModel } from 'vs/workbench/contrib/notebook/browser/
import { CellEditState, CellFindMatchWithIndex, CellWebviewFindMatch, ICellViewModel } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
import { NotebookEditorWidget } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/services/notebookEditorService';
import { NotebookCellsChangeType } from 'vs/workbench/contrib/notebook/common/notebookCommon';
import { NotebookCellsChangeType, NotebookData } from 'vs/workbench/contrib/notebook/common/notebookCommon';
import { INotebookService, SimpleNotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookService';
import { IReplaceService } from 'vs/workbench/contrib/search/browser/replace';
import { ICellMatch, IFileMatchWithCells, contentMatchesToTextSearchMatches, isIFileMatchWithCells, webviewMatchesToTextSearchMatches } from 'vs/workbench/contrib/search/browser/searchNotebookHelpers';
import { CellSearchModel, ICellMatch, IFileMatchWithCells, contentMatchesToTextSearchMatches, isIFileMatchWithCells, webviewMatchesToTextSearchMatches } from 'vs/workbench/contrib/search/browser/searchNotebookHelpers';
import { ReplacePattern } from 'vs/workbench/services/search/common/replace';
import { IFileMatch, IPatternInfo, ISearchComplete, ISearchConfigurationProperties, ISearchProgressItem, ISearchRange, ISearchService, ITextQuery, ITextSearchContext, ITextSearchMatch, ITextSearchPreviewOptions, ITextSearchResult, ITextSearchStats, OneLineRange, QueryType, resultIsMatch, SearchCompletionExitCode, SearchSortOrder } from 'vs/workbench/services/search/common/search';
import { IFileMatch, IFileQuery, IPatternInfo, ISearchComplete, ISearchConfigurationProperties, ISearchProgressItem, ISearchRange, ISearchService, ITextQuery, ITextSearchContext, ITextSearchMatch, ITextSearchPreviewOptions, ITextSearchResult, ITextSearchStats, OneLineRange, QueryType, resultIsMatch, SearchCompletionExitCode, SearchSortOrder } from 'vs/workbench/services/search/common/search';
import { addContextToEditorMatches, editorMatchesToTextSearchResults } from 'vs/workbench/services/search/common/searchHelpers';

export class Match {
Expand Down Expand Up @@ -187,7 +189,7 @@ export class CellMatch {

constructor(
private readonly _parent: FileMatch,
private readonly _cell: ICellViewModel,
private readonly _cell: ICellViewModel | CellSearchModel,
private readonly _cellIndex: number,
) {

Expand Down Expand Up @@ -231,6 +233,9 @@ export class CellMatch {
}

public addContext(textSearchMatches: ITextSearchMatch[]) {
if (this.cell instanceof CellSearchModel) {
return;
}
this.cell.resolveTextModel().then((textModel) => {
const textResultsWithContext = addContextToEditorMatches(textSearchMatches, textModel, this.parent.parent().query!);
const contexts = textResultsWithContext.filter((result => !resultIsMatch(result)) as ((a: any) => a is ITextSearchContext));
Expand Down Expand Up @@ -260,7 +265,7 @@ export class CellMatch {
return this._cellIndex;
}

get cell(): ICellViewModel {
get cell(): ICellViewModel | CellSearchModel {
return this._cell;
}

Expand Down Expand Up @@ -812,7 +817,7 @@ export class FileMatch extends Disposable implements IFileMatch {
}

private async highlightCurrentFindMatchDecoration(match: MatchInNotebook): Promise<number | null> {
if (!this._findMatchDecorationModel) {
if (!this._findMatchDecorationModel || match.cell instanceof CellSearchModel) {
return null;
}
if (match.webviewIndex === undefined) {
Expand All @@ -823,7 +828,7 @@ export class FileMatch extends Disposable implements IFileMatch {
}

private revealCellRange(match: MatchInNotebook, outputOffset: number | null) {
if (!this._notebookEditorWidget) {
if (!this._notebookEditorWidget || match.cell instanceof CellSearchModel) {
return;
}
if (match.webviewIndex !== undefined) {
Expand Down Expand Up @@ -1906,6 +1911,9 @@ export class SearchModel extends Disposable {
@IInstantiationService private readonly instantiationService: IInstantiationService,
@IUriIdentityService private readonly uriIdentityService: IUriIdentityService,
@INotebookEditorService private readonly notebookEditorService: INotebookEditorService,
@INotebookService private readonly notebookService: INotebookService,
@IFileService private readonly fileService: IFileService,
@ILogService private readonly logService: ILogService,
) {
super();
this._searchResult = this.instantiationService.createInstance(SearchResult, this);
Expand Down Expand Up @@ -1946,6 +1954,74 @@ export class SearchModel extends Disposable {
get searchResult(): SearchResult {
return this._searchResult;
}
private async getClosedNotebookResults(textQuery: ITextQuery, token: CancellationToken): Promise<{ results: ResourceMap<IFileMatchWithCells | null>; limitHit: boolean }> {
const results = new ResourceMap<IFileMatchWithCells | null>(uri => this.uriIdentityService.extUri.getComparisonKey(uri));
const query: IFileQuery = {
type: QueryType.File,
filePattern: '**/*.ipynb',
folderQueries: textQuery.folderQueries
};

const searchComplete = await this.searchService.fileSearch(
query,
CancellationToken.None
);
for (const fileMatch of searchComplete.results) {
const cellMatches: ICellMatch[] = [];
const uri = fileMatch.resource;
const content = await this.fileService.readFileStream(uri);
try {
const info = await this.notebookService.withNotebookDataProvider('jupyter-notebook');
if (!(info instanceof SimpleNotebookProviderInfo)) {
throw new Error('CANNOT open file notebook with this provider');
}

let _data: NotebookData = {
metadata: {},
cells: []
};
if (uri.scheme !== Schemas.vscodeInteractive) {
const bytes = await streamToBuffer(content.value);
_data = await info.serializer.dataToNotebook(bytes);
}


_data.cells.forEach((cell, index) => {
const source = cell.source;
const target = textQuery.contentPattern.pattern;

const cellModel = new CellSearchModel(source, uri, index);
const matches = cellModel.find(target);
if (matches.length > 0) {
const cellMatch: ICellMatch = {
cell: cellModel,
index: index,
contentResults: contentMatchesToTextSearchMatches(matches, cellModel),
webviewResults: []
};
cellMatches.push(cellMatch);
}
});

if (cellMatches.length > 0) {
const fileMatch: IFileMatchWithCells = {
resource: uri, cellResults: cellMatches
};
results.set(uri, fileMatch);
}

} catch (e) {
this.logService.info('error: ' + e);
continue;
}
}

return {
results: results,
limitHit: false
};
}


private async getLocalNotebookResults(query: ITextQuery, token: CancellationToken): Promise<{ results: ResourceMap<IFileMatchWithCells | null>; limitHit: boolean }> {
const localResults = new ResourceMap<IFileMatchWithCells | null>(uri => this.uriIdentityService.extUri.getComparisonKey(uri));
Expand Down Expand Up @@ -2006,16 +2082,18 @@ export class SearchModel extends Disposable {

async notebookSearch(query: ITextQuery, token: CancellationToken, onProgress?: (result: ISearchProgressItem) => void): Promise<{ completeData: ISearchComplete; scannedFiles: ResourceSet }> {
const localResults = await this.getLocalNotebookResults(query, token);
const closedResults = await this.getClosedNotebookResults(query, token);

if (onProgress) {
arrays.coalesce([...localResults.results.values()]).forEach(onProgress);
arrays.coalesce([...localResults.results.values(), ...closedResults.results.values()]).forEach(onProgress);
}
return {
completeData: {
messages: [],
limitHit: localResults.limitHit,
results: arrays.coalesce([...localResults.results.values()]),
results: arrays.coalesce([...localResults.results.values(), ...closedResults.results.values()]),
},
scannedFiles: new ResourceSet([...localResults.results.keys()], uri => this.uriIdentityService.extUri.getComparisonKey(uri))
scannedFiles: new ResourceSet([...localResults.results.keys(), ...closedResults.results.keys()], uri => this.uriIdentityService.extUri.getComparisonKey(uri))
};
}

Expand Down
63 changes: 60 additions & 3 deletions src/vs/workbench/contrib/search/browser/searchNotebookHelpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,21 @@
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/

import { FindMatch } from 'vs/editor/common/model';
import { DefaultEndOfLine, FindMatch, IReadonlyTextBuffer } from 'vs/editor/common/model';
import { CellWebviewFindMatch, ICellViewModel } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
import { IFileMatch, ITextSearchMatch, TextSearchMatch } from 'vs/workbench/services/search/common/search';
import { Range } from 'vs/editor/common/core/range';
import { PieceTreeTextBufferBuilder } from 'vs/editor/common/model/pieceTreeTextBuffer/pieceTreeTextBufferBuilder';
import { SearchParams } from 'vs/editor/common/model/textModelSearch';
import { Disposable } from 'vs/base/common/lifecycle';
import { URI } from 'vs/base/common/uri';

export interface IFileMatchWithCells extends IFileMatch {
cellResults: ICellMatch[];
}

export interface ICellMatch {
cell: ICellViewModel;
cell: ICellViewModel | CellSearchModel;
index: number;
contentResults: ITextSearchMatch[];
webviewResults: ITextSearchMatch[];
Expand All @@ -24,7 +28,7 @@ export function isIFileMatchWithCells(object: IFileMatch): object is IFileMatchW

// to text search results

export function contentMatchesToTextSearchMatches(contentMatches: FindMatch[], cell: ICellViewModel): ITextSearchMatch[] {
export function contentMatchesToTextSearchMatches(contentMatches: FindMatch[], cell: ICellViewModel | CellSearchModel): ITextSearchMatch[] {
let previousEndLine = -1;
const contextGroupings: FindMatch[][] = [];
let currentContextGrouping: FindMatch[] = [];
Expand Down Expand Up @@ -73,4 +77,57 @@ export function webviewMatchesToTextSearchMatches(webviewMatches: CellWebviewFin
).filter((e): e is ITextSearchMatch => !!e);
}

// experimental

export class CellSearchModel extends Disposable {
constructor(readonly _source: string, private _uri: URI, private _cellIndex: number) {
// need cell index
super();
}

get id() {
return `${this._uri.toString()}#${this._cellIndex}`;
}

get uri() {
return this._uri;
}

public getFullModelRange(): Range {
const lineCount = this.textBuffer.getLineCount();
return new Range(1, 1, lineCount, this.getLineMaxColumn(lineCount));
}

public getLineMaxColumn(lineNumber: number): number {
if (lineNumber < 1 || lineNumber > this.textBuffer.getLineCount()) {
throw new Error('Illegal value for lineNumber');
}
return this.textBuffer.getLineLength(lineNumber) + 1;
}

private _textBuffer!: IReadonlyTextBuffer;
get textBuffer() {
if (this._textBuffer) {
return this._textBuffer;
}

const builder = new PieceTreeTextBufferBuilder();
builder.acceptChunk(this._source);
const bufferFactory = builder.finish(true);
const { textBuffer, disposable } = bufferFactory.create(DefaultEndOfLine.LF);
this._textBuffer = textBuffer;
this._register(disposable);

return this._textBuffer;
}

find(target: string): FindMatch[] {
const searchParams = new SearchParams(target, false, false, null);
const searchData = searchParams.parseSearchRequest();
if (!searchData) {
return [];
}
const fullRange = this.getFullModelRange();
return this.textBuffer.findMatchesLineByLine(fullRange, searchData, true, 5000);
}
}

0 comments on commit 2abfedc

Please sign in to comment.