jupyter notebook variable provider improvements (#15018)

* failing test

* passing tests

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

@@ -0,0 +1,171 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { assert } from 'chai';
+import { JupyterVariablesProvider } from './JupyterVariablesProvider';
+import { NotebookDocument, CancellationTokenSource, VariablesResult, Variable } from 'vscode';
+import { mock, instance, when, anything, verify, objectContaining } from 'ts-mockito';
+import { IKernelProvider, IKernel } from '../types';
+import { IJupyterVariables, IVariableDescription } from './types';
+
+suite('JupyterVariablesProvider', () => {
+    let variables: IJupyterVariables;
+    let kernelProvider: IKernelProvider;
+    let provider: JupyterVariablesProvider;
+    const notebook = mock<NotebookDocument>();
+    const cancellationToken = new CancellationTokenSource().token;
+
+    const objectVariable: IVariableDescription = {
+        name: 'myObject',
+        value: '...',
+        root: 'myObject',
+        properties: ['myList'],
+        propertyChain: []
+    };
+
+    const listVariable: IVariableDescription = {
+        name: 'myList',
+        value: '[...]',
+        count: 3,
+        root: 'myObject',
+        propertyChain: ['myList']
+    };
+
+    function createListItem(index: number): IVariableDescription {
+        return {
+            name: index.toString(),
+            value: `value${index}`,
+            count: index,
+            root: 'myObject',
+            propertyChain: ['myList', index]
+        };
+    }
+
+    async function provideVariables(parent: Variable | undefined) {
+        const results: VariablesResult[] = [];
+        for await (const result of provider.provideVariables(
+            instance(notebook),
+            parent,
+            1, // Named
+            0,
+            cancellationToken
+        )) {
+            results.push(result);
+        }
+        return results;
+    }
+
+    const listVariableItems = [0, 1, 2].map(createListItem);
+
+    setup(() => {
+        variables = mock<IJupyterVariables>();
+        kernelProvider = mock<IKernelProvider>();
+        provider = new JupyterVariablesProvider(instance(variables), instance(kernelProvider));
+    });
+
+    test('provideVariables without parent should yield variables', async () => {
+        const kernel = mock<IKernel>();
+
+        when(kernelProvider.get(anything())).thenReturn(instance(kernel));
+        when(variables.getAllVariableDiscriptions(anything(), undefined, anything())).thenReturn(
+            Promise.resolve([objectVariable])
+        );
+        when(kernelProvider.getKernelExecution(anything())).thenReturn({ executionCount: 1 } as any);
+
+        const results = await provideVariables(undefined);
+
+        assert.isNotEmpty(results);
+        assert.equal(results.length, 1);
+        assert.equal(results[0].variable.name, 'myObject');
+    });
+
+    test('provideVariables with a parent should call get children correctly', async () => {
+        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);
+
+        // 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);
+
+        assert.equal(listResult.variable.name, 'myList');
+        assert.isNotEmpty(listItems);
+        assert.equal(listItems.length, 3);
+        listItems.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 = {
+            name: 'myInt',
+            value: '1',
+            root: '',
+            propertyChain: []
+        };
+
+        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);
+
+        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');
+        assert.equal(second[0].variable.value, '2');
+    });
+
+    test('Getting variables again with same execution count should not make another call', async () => {
+        const kernel = mock<IKernel>();
+        const intVariable: IVariableDescription = {
+            name: 'myInt',
+            value: '1',
+            root: '',
+            propertyChain: []
+        };
+
+        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);
+
+        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();
+    });
+});

src/kernels/variables/JupyterVariablesProvider.ts

@@ -12,16 +12,19 @@ import {
 } from 'vscode';
 import { IJupyterVariables, IVariableDescription } from './types';
 import { IKernel, IKernelProvider } from '../types';
+import { VariableResultCache } from './variableResultCache';
 
 export class JupyterVariablesProvider implements NotebookVariableProvider {
+    private variableResultCache = new VariableResultCache();
+
+    _onDidChangeVariables = new EventEmitter<NotebookDocument>();
+    onDidChangeVariables = this._onDidChangeVariables.event;
+
     constructor(
         private readonly variables: IJupyterVariables,
         private readonly kernelProvider: IKernelProvider
     ) {}
 
-    _onDidChangeVariables = new EventEmitter<NotebookDocument>();
-    onDidChangeVariables = this._onDidChangeVariables.event;
-
     async *provideVariables(
         notebook: NotebookDocument,
         parent: Variable | undefined,
@@ -37,34 +40,51 @@ export class JupyterVariablesProvider implements NotebookVariableProvider {
             return;
         }
 
-        if (parent) {
-            if ('getChildren' in parent && typeof parent.getChildren === 'function') {
-                const variables = await parent.getChildren(start);
-                for (const variable of variables) {
-                    yield this.createVariableResult(variable, kernel);
+        const executionCount = this.kernelProvider.getKernelExecution(kernel).executionCount;
+
+        const cacheKey = this.variableResultCache.getCacheKey(notebook.uri.toString(), parent);
+        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);
+                results = variables.map((variable) => this.createVariableResult(variable, kernel));
             }
-        } else {
-            const variables = await this.variables.getAllVariableDiscriptions(kernel, undefined);
+        }
 
-            for (const variable of variables) {
-                yield this.createVariableResult(variable, kernel);
-            }
+        if (!results) {
+            return;
+        }
+
+        this.variableResultCache.setResults(executionCount, cacheKey, results);
+
+        for (const result of results) {
+            yield result;
         }
     }
 
     private createVariableResult(result: IVariableDescription, kernel: IKernel): VariablesResult {
         const hasNamedChildren = !!result.properties;
         const indexedChildrenCount = result.count ?? 0;
         const variable = {
-            getChildren: (start: number) => this.getChildren(variable, start, kernel),
+            getChildren: (start: number, token: CancellationToken) => this.getChildren(variable, start, kernel, token),
             ...result
         } as Variable;
         return { variable, hasNamedChildren, indexedChildrenCount };
     }
 
-    async getChildren(variable: Variable, _start: number, kernel: IKernel): Promise<IVariableDescription[]> {
+    async getChildren(
+        variable: Variable,
+        _start: number,
+        kernel: IKernel,
+        token: CancellationToken
+    ): Promise<IVariableDescription[]> {
         const parent = variable as IVariableDescription;
-        return await this.variables.getAllVariableDiscriptions(kernel, parent);
+        return await this.variables.getAllVariableDiscriptions(kernel, parent, token);
     }
 }

src/kernels/variables/pythonVariableRequester.ts

@@ -153,7 +153,7 @@ export class PythonVariablesRequester implements IKernelVariableRequester {
         parent: IVariableDescription | undefined,
         token?: CancellationToken
     ): Promise<IVariableDescription[]> {
-        if (!kernel.session || token?.isCancellationRequested) {
+        if (!kernel.session) {
             return [];
         }
 

src/kernels/variables/variableResultCache.ts

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { Variable, VariablesResult } from 'vscode';
+import { IVariableDescription } from './types';
+
+export class VariableResultCache {
+    private cache = new Map<string, VariablesResult[]>();
+    private executionCount = 0;
+
+    getCacheKey(notebookUri: string, parent: Variable | undefined): string {
+        let parentKey = '';
+        const parentDescription = parent as IVariableDescription;
+        if (parentDescription) {
+            parentKey = `${parentDescription.name}.${parentDescription.propertyChain.join('.')}`;
+        }
+        return `${notebookUri}:${parentKey}`;
+    }
+
+    getResults(executionCount: number, cacheKey: string): VariablesResult[] | undefined {
+        if (this.executionCount !== executionCount) {
+            this.cache.clear();
+            this.executionCount = executionCount;
+        }
+
+        return this.cache.get(cacheKey);
+    }
+
+    setResults(executionCount: number, cacheKey: string, results: VariablesResult[]) {
+        if (this.executionCount < executionCount) {
+            this.cache.clear();
+            this.executionCount = executionCount;
+        } else if (this.executionCount > executionCount) {
+            // old results, don't cache
+            return;
+        }
+
+        this.cache.set(cacheKey, results);
+    }
+}