Skip to content

Commit

Permalink
return indexed children on demand (#15023)
Browse files Browse the repository at this point in the history 
* limit collection size returned

* get indexed children in pages
  • Loading branch information
@amunger
amunger authored Jan 17, 2024
1 parent 2e3a664 commit f2e694e
Show file tree
Hide file tree
Showing 10 changed files with 169 additions and 71 deletions.
9 changes: 7 additions & 2 deletions pythonFiles/vscode_datascience_helpers/getVariableInfo/vscodeGetVariableInfo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ def _VSCODE_getVariable(what_to_get, is_debugging, *args):

maxStringLength = 50
collectionTypes = ["list", "tuple", "set"]
arrayPageSize = 50

def truncateString(string):
if _VSCODE_builtins.len(string) > maxStringLength:
Expand Down Expand Up @@ -95,7 +96,7 @@ def _VSCODE_getAllVariableDescriptions(varNames):
return _VSCODE_builtins.print(_VSCODE_json.dumps(variables))

### Get info on children of a variable reached through the given property chain
def _VSCODE_getAllChildrenDescriptions(rootVarName, propertyChain=[]):
def _VSCODE_getAllChildrenDescriptions(rootVarName, propertyChain, startIndex):
root = globals()[rootVarName]
if root is None:
return []
Expand All @@ -108,14 +109,18 @@ def _VSCODE_getAllChildrenDescriptions(rootVarName, propertyChain=[]):
parentInfo = getVariableDescription(parent)
if "count" in parentInfo:
if parentInfo["count"] > 0:
lastItem = _VSCODE_builtins.min(
parentInfo["count"], startIndex + arrayPageSize
)
range = _VSCODE_builtins.range(startIndex, lastItem)
children = [
{
**getVariableDescription(getChildProperty(parent, [i])),
"name": str(i),
"root": rootVarName,
"propertyChain": propertyChain + [i],
}
for i in _VSCODE_builtins.range(_VSCODE_builtins.len(parent))
for i in range
]
elif "properties" in parentInfo:
children = [
Expand Down
141 changes: 101 additions & 40 deletions src/kernels/variables/JupyterVariableProvider.unit.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,22 +41,30 @@ suite('JupyterVariablesProvider', () => {
};
}

async function provideVariables(parent: Variable | undefined) {
function setVariablesForParent(
parent: IVariableDescription | undefined,
result: IVariableDescription[],
updated?: IVariableDescription[],
startIndex?: number
) {
const whenCallHappens = when(
variables.getAllVariableDiscriptions(anything(), parent, startIndex ?? anything(), anything())
);

whenCallHappens.thenReturn(Promise.resolve(result));
if (updated) {
whenCallHappens.thenReturn(Promise.resolve(updated));
}
}

async function provideVariables(parent: Variable | undefined, kind = 1) {
const results: VariablesResult[] = [];
for await (const result of provider.provideVariables(
instance(notebook),
parent,
1, // Named
0,
cancellationToken
)) {
for await (const result of provider.provideVariables(instance(notebook), parent, kind, 0, cancellationToken)) {
results.push(result);
}
return results;
}

const listVariableItems = [0, 1, 2].map(createListItem);

setup(() => {
variables = mock<IJupyterVariables>();
kernelProvider = mock<IKernelProvider>();
Expand All @@ -67,9 +75,7 @@ suite('JupyterVariablesProvider', () => {
const kernel = mock<IKernel>();

when(kernelProvider.get(anything())).thenReturn(instance(kernel));
when(variables.getAllVariableDiscriptions(anything(), undefined, anything())).thenReturn(
Promise.resolve([objectVariable])
);
setVariablesForParent(undefined, [objectVariable]);
when(kernelProvider.getKernelExecution(anything())).thenReturn({ executionCount: 1 } as any);

const results = await provideVariables(undefined);
Expand All @@ -83,29 +89,17 @@ suite('JupyterVariablesProvider', () => {
const kernel = mock<IKernel>();

when(kernelProvider.get(anything())).thenReturn(instance(kernel));
when(variables.getAllVariableDiscriptions(anything(), undefined, anything())).thenReturn(
Promise.resolve([objectVariable])
);
when(
variables.getAllVariableDiscriptions(
anything(),
objectContaining({ root: 'myObject', propertyChain: [] }),
anything()
)
).thenReturn(Promise.resolve([listVariable]));
when(
variables.getAllVariableDiscriptions(
anything(),
objectContaining({ root: 'myObject', propertyChain: ['myList'] }),
anything()
)
).thenReturn(Promise.resolve(listVariableItems));
when(kernelProvider.getKernelExecution(anything())).thenReturn({ executionCount: 1 } as any);

const listVariableItems = [0, 1, 2].map(createListItem);
setVariablesForParent(undefined, [objectVariable]);
setVariablesForParent(objectContaining({ root: 'myObject', propertyChain: [] }), [listVariable]);
setVariablesForParent(objectContaining({ root: 'myObject', propertyChain: ['myList'] }), listVariableItems);

// pass each the result as the parent in the next call
let rootVariable = (await provideVariables(undefined))[0];
const listResult = (await provideVariables(rootVariable!.variable))[0];
const listItems = await provideVariables(listResult!.variable);
const listItems = await provideVariables(listResult!.variable, 2);

assert.equal(listResult.variable.name, 'myList');
assert.isNotEmpty(listItems);
Expand All @@ -116,6 +110,77 @@ suite('JupyterVariablesProvider', () => {
});
});

test('All indexed variables should be returned when requested', async () => {
const kernel = mock<IKernel>();

const listVariable: IVariableDescription = {
name: 'myList',
value: '[...]',
count: 6,
root: 'myList',
propertyChain: []
};

when(kernelProvider.get(anything())).thenReturn(instance(kernel));
when(kernelProvider.getKernelExecution(anything())).thenReturn({ executionCount: 1 } as any);

const firstPage = [0, 1, 2].map(createListItem);
const secondPage = [3, 4, 5].map(createListItem);
setVariablesForParent(undefined, [listVariable]);
setVariablesForParent(objectContaining({ root: 'myList', propertyChain: [] }), firstPage, undefined, 0);
setVariablesForParent(
objectContaining({ root: 'myList', propertyChain: [] }),
secondPage,
undefined,
firstPage.length
);

const rootVariable = (await provideVariables(undefined))[0];
const listItemResult = await provideVariables(rootVariable!.variable, 2);

assert.equal(listItemResult.length, 6, 'full list of items should be returned');
listItemResult.forEach((item, index) => {
assert.equal(item.variable.name, index.toString());
assert.equal(item.variable.value, `value${index}`);
});
});

test('Getting less indexed items than the specified count is handled', async () => {
const kernel = mock<IKernel>();

const listVariable: IVariableDescription = {
name: 'myList',
value: '[...]',
count: 6,
root: 'myList',
propertyChain: []
};

when(kernelProvider.get(anything())).thenReturn(instance(kernel));
when(kernelProvider.getKernelExecution(anything())).thenReturn({ executionCount: 1 } as any);

const firstPage = [0, 1, 2].map(createListItem);
const secondPage = [3, 4].map(createListItem);
setVariablesForParent(undefined, [listVariable]);
setVariablesForParent(objectContaining({ root: 'myList', propertyChain: [] }), firstPage, undefined, 0);
setVariablesForParent(
objectContaining({ root: 'myList', propertyChain: [] }),
secondPage,
undefined,
firstPage.length
);
setVariablesForParent(objectContaining({ root: 'myList', propertyChain: [] }), [], undefined, 5);

const rootVariable = (await provideVariables(undefined))[0];
const listItemResult = await provideVariables(rootVariable!.variable, 2);

assert.equal(listItemResult.length, 5);
listItemResult.forEach((item, index) => {
assert.equal(item.variable.name, index.toString());
assert.equal(item.variable.value, `value${index}`);
});
});

test('Getting variables again with new execution count should get updated variables', async () => {
const kernel = mock<IKernel>();
const intVariable: IVariableDescription = {
Expand All @@ -126,14 +191,12 @@ suite('JupyterVariablesProvider', () => {
};

when(kernelProvider.get(anything())).thenReturn(instance(kernel));
when(variables.getAllVariableDiscriptions(anything(), undefined, anything()))
.thenReturn(Promise.resolve([intVariable]))
.thenReturn(Promise.resolve([{ ...intVariable, value: '2' }]));

when(kernelProvider.getKernelExecution(anything()))
.thenReturn({ executionCount: 1 } as any)
.thenReturn({ executionCount: 2 } as any);

setVariablesForParent(undefined, [intVariable], [{ ...intVariable, value: '2' }]);

const first = await provideVariables(undefined);
const second = await provideVariables(undefined);

Expand All @@ -153,19 +216,17 @@ suite('JupyterVariablesProvider', () => {
};

when(kernelProvider.get(anything())).thenReturn(instance(kernel));
when(variables.getAllVariableDiscriptions(anything(), undefined, anything())).thenReturn(
Promise.resolve([intVariable])
);

when(kernelProvider.getKernelExecution(anything())).thenReturn({ executionCount: 1 } as any);

setVariablesForParent(undefined, [intVariable]);

const first = await provideVariables(undefined);
const second = await provideVariables(undefined);

assert.equal(first.length, 1);
assert.equal(second.length, 1);
assert.equal(first[0].variable.value, '1');

verify(variables.getAllVariableDiscriptions(anything(), undefined, anything())).once();
verify(variables.getAllVariableDiscriptions(anything(), anything(), anything(), anything())).once();
});
});
60 changes: 41 additions & 19 deletions src/kernels/variables/JupyterVariablesProvider.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ export class JupyterVariablesProvider implements NotebookVariableProvider {
async *provideVariables(
notebook: NotebookDocument,
parent: Variable | undefined,
_kind: NotebookVariablesRequestKind,
kind: NotebookVariablesRequestKind,
start: number,
token: CancellationToken
): AsyncIterable<VariablesResult> {
Expand All @@ -42,29 +42,51 @@ export class JupyterVariablesProvider implements NotebookVariableProvider {

const executionCount = this.kernelProvider.getKernelExecution(kernel).executionCount;

const cacheKey = this.variableResultCache.getCacheKey(notebook.uri.toString(), parent);
const cacheKey = this.variableResultCache.getCacheKey(notebook.uri.toString(), parent, start);
let results = this.variableResultCache.getResults(executionCount, cacheKey);

if (!results) {
if (parent) {
if ('getChildren' in parent && typeof parent.getChildren === 'function') {
const variables = (await parent.getChildren(start, token)) as IVariableDescription[];
results = variables.map((variable) => this.createVariableResult(variable, kernel));
}
} else {
const variables = await this.variables.getAllVariableDiscriptions(kernel, undefined, token);
if (parent) {
const parentDescription = parent as IVariableDescription;
if (!results && parentDescription.getChildren) {
const variables = await parentDescription.getChildren(start, token);
results = variables.map((variable) => this.createVariableResult(variable, kernel));
this.variableResultCache.setResults(executionCount, cacheKey, results);
} else {
// no cached results and no way to get children, so return empty
return;
}
}

if (!results) {
return;
}
for (const result of results) {
yield result;
}

this.variableResultCache.setResults(executionCount, cacheKey, results);
// check if we have more indexed children to return
if (
kind === 2 &&
parentDescription.count &&
results.length > 0 &&
parentDescription.count > start + results.length
) {
for await (const result of this.provideVariables(
notebook,
parent,
kind,
start + results.length,
token
)) {
yield result;
}
}
} else {
if (!results) {
const variables = await this.variables.getAllVariableDiscriptions(kernel, undefined, start, token);
results = variables.map((variable) => this.createVariableResult(variable, kernel));
this.variableResultCache.setResults(executionCount, cacheKey, results);
}

for (const result of results) {
yield result;
for (const result of results) {
yield result;
}
}
}

Expand All @@ -80,11 +102,11 @@ export class JupyterVariablesProvider implements NotebookVariableProvider {

async getChildren(
variable: Variable,
_start: number,
start: number,
kernel: IKernel,
token: CancellationToken
): Promise<IVariableDescription[]> {
const parent = variable as IVariableDescription;
return await this.variables.getAllVariableDiscriptions(kernel, parent, token);
return await this.variables.getAllVariableDiscriptions(kernel, parent, start, token);
}
}
5 changes: 3 additions & 2 deletions src/kernels/variables/jupyterVariables.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,10 @@ export class JupyterVariables implements IJupyterVariables {
getAllVariableDiscriptions(
kernel: IKernel,
parent: IVariableDescription | undefined,
token?: CancellationToken
startIndex: number,
token: CancellationToken
): Promise<IVariableDescription[]> {
return this.variableHandler.getAllVariableDiscriptions(kernel, parent, token);
return this.variableHandler.getAllVariableDiscriptions(kernel, parent, startIndex, token);
}

@capturePerfTelemetry(Telemetry.VariableExplorerFetchTime)
Expand Down
5 changes: 3 additions & 2 deletions src/kernels/variables/kernelVariables.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,12 +75,13 @@ export class KernelVariables implements IJupyterVariables {
public async getAllVariableDiscriptions(
kernel: IKernel,
parent: IVariableDescription | undefined,
token?: CancellationToken
startIndex: number,
token: CancellationToken
): Promise<IVariableDescription[]> {
const languageId = getKernelConnectionLanguage(kernel.kernelConnectionMetadata) || PYTHON_LANGUAGE;
const variableRequester = this.variableRequesters.get(languageId);
if (variableRequester) {
return variableRequester.getAllVariableDiscriptions(kernel, parent, token);
return variableRequester.getAllVariableDiscriptions(kernel, parent, startIndex, token);
}
return [];
}
Expand Down
4 changes: 3 additions & 1 deletion src/kernels/variables/pythonVariableRequester.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,7 @@ export class PythonVariablesRequester implements IKernelVariableRequester {
public async getAllVariableDiscriptions(
kernel: IKernel,
parent: IVariableDescription | undefined,
startIndex: number,
token?: CancellationToken
): Promise<IVariableDescription[]> {
if (!kernel.session) {
Expand All @@ -160,7 +161,8 @@ export class PythonVariablesRequester implements IKernelVariableRequester {
const { code, cleanupCode, initializeCode } =
await this.varScriptGenerator.generateCodeToGetAllVariableDescriptions({
isDebugging: false,
parent
parent,
startIndex
});

const results = await safeExecuteSilently(
Expand Down
Loading

0 comments on commit f2e694e

Please sign in to comment.