return indexed children on demand

pythonFiles/vscode_datascience_helpers/getVariableInfo/vscodeGetVariableInfo.py

@@ -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:
@@ -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 []
@@ -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 = [

src/kernels/variables/JupyterVariableProvider.unit.test.ts

@@ -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>();
@@ -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);
@@ -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);
@@ -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 = {
@@ -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);
 
@@ -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();
     });
 });

src/kernels/variables/JupyterVariablesProvider.ts

@@ -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> {
@@ -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;
+            }
         }
     }
 
@@ -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);
     }
 }

src/kernels/variables/jupyterVariables.ts

@@ -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)

src/kernels/variables/kernelVariables.ts

@@ -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 [];
     }

src/kernels/variables/pythonVariableRequester.ts

@@ -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) {
@@ -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(