Add extra code to compute shape for python objects. (#11200)

* Add extra code to compute shape for python objects.

* Fix autocomplete tests. Not sure how these are passing in master

* May as well skip list if not necessary

* Fix native editor intellisense unit tests

Author: rchiodo

news/2 Fixes/10825.md

@@ -0,0 +1 @@
+Fix problem with shape not being computed for some types in the variable explorer.

pythonFiles/vscode_datascience_helpers/getJupyterVariableShape.py

@@ -0,0 +1,35 @@
+# Query Jupyter server for defined variables list
+# Tested on 2.7 and 3.6
+import json as _VSCODE_json
+
+# In IJupyterVariables.getValue this '_VSCode_JupyterTestValue' will be replaced with the json stringified value of the target variable
+# Indexes off of _VSCODE_targetVariable need to index types that are part of IJupyterVariable
+_VSCODE_targetVariable = _VSCODE_json.loads("""_VSCode_JupyterTestValue""")
+
+_VSCODE_evalResult = eval(_VSCODE_targetVariable["name"])
+
+# Find shape and count if available
+if hasattr(_VSCODE_evalResult, "shape"):
+    try:
+        # Get a bit more restrictive with exactly what we want to count as a shape, since anything can define it
+        if isinstance(_VSCODE_evalResult.shape, tuple):
+            _VSCODE_shapeStr = str(_VSCODE_evalResult.shape)
+            if (
+                len(_VSCODE_shapeStr) >= 3
+                and _VSCODE_shapeStr[0] == "("
+                and _VSCODE_shapeStr[-1] == ")"
+                and "," in _VSCODE_shapeStr
+            ):
+                _VSCODE_targetVariable["shape"] = _VSCODE_shapeStr
+            del _VSCODE_shapeStr
+    except TypeError:
+        pass
+
+if hasattr(_VSCODE_evalResult, "__len__"):
+    try:
+        _VSCODE_targetVariable["count"] = len(_VSCODE_evalResult)
+    except TypeError:
+        pass
+
+print(_VSCODE_json.dumps(_VSCODE_targetVariable))
+del _VSCODE_json

src/client/common/process/pythonDaemon.ts

@@ -49,6 +49,7 @@ export class PythonDaemonExecutionService implements IPythonDaemonExecutionServi
     public get isAlive(): boolean {
         return this.connectionClosedMessage === '';
     }
+    private disposed = false;
     constructor(
         protected readonly pythonExecutionService: IPythonExecutionService,
         protected readonly pythonPath: string,
@@ -64,6 +65,7 @@ export class PythonDaemonExecutionService implements IPythonDaemonExecutionServi
     }
     public dispose() {
         try {
+            this.disposed = true;
             // The daemon should die as a result of this.
             this.connection.sendNotification(new NotificationType('exit'));
             this.proc.kill();
@@ -382,11 +384,13 @@ export class PythonDaemonExecutionService implements IPythonDaemonExecutionServi
     private monitorConnection() {
         // tslint:disable-next-line: no-any
         const logConnectionStatus = (msg: string, ex?: any) => {
-            this.connectionClosedMessage += msg + (ex ? `, With Error: ${util.format(ex)}` : '');
-            this.connectionClosedDeferred.reject(new ConnectionClosedError(this.connectionClosedMessage));
-            traceWarning(msg);
-            if (ex) {
-                traceError('Connection errored', ex);
+            if (!this.disposed) {
+                this.connectionClosedMessage += msg + (ex ? `, With Error: ${util.format(ex)}` : '');
+                this.connectionClosedDeferred.reject(new ConnectionClosedError(this.connectionClosedMessage));
+                traceWarning(msg);
+                if (ex) {
+                    traceError('Connection errored', ex);
+                }
             }
         };
         this.disposables.push(this.connection.onClose(() => logConnectionStatus('Daemon Connection Closed')));

src/client/datascience/interactive-common/intellisense/intellisenseDocument.ts

@@ -207,8 +207,11 @@ export class IntellisenseDocument implements TextDocument {
         };
     }
 
-    public loadAllCells(cells: { code: string; id: string }[]): TextDocumentContentChangeEvent[] {
-        if (!this.inEditMode) {
+    public loadAllCells(
+        cells: { code: string; id: string }[],
+        notebookType: 'interactive' | 'native'
+    ): TextDocumentContentChangeEvent[] {
+        if (!this.inEditMode && notebookType === 'native') {
             this.inEditMode = true;
             return this.reloadCells(cells);
         }
@@ -593,7 +596,9 @@ export class IntellisenseDocument implements TextDocument {
         }
 
         // in interactive window
-        return this._cellRanges[this._cellRanges.length - 1].start;
+        return this._cellRanges && this._cellRanges.length > 0
+            ? this._cellRanges[this._cellRanges.length - 1].start
+            : 0;
     }
 
     private getLineFromOffset(offset: number) {

src/client/datascience/interactive-common/intellisense/intellisenseProvider.ts

@@ -68,6 +68,7 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
     }>();
     private cancellationSources: Map<string, CancellationTokenSource> = new Map<string, CancellationTokenSource>();
     private notebookIdentity: Uri | undefined;
+    private notebookType: 'interactive' | 'native' = 'interactive';
     private potentialResource: Uri | undefined;
     private sentOpenDocument: boolean = false;
     private languageServer: ILanguageServer | undefined;
@@ -688,7 +689,8 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
                             code: concatMultilineStringInput(cell.data.source),
                             id: cell.id
                         };
-                    })
+                    }),
+                this.notebookType
             );
 
             await this.handleChanges(document, changes);
@@ -708,6 +710,7 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
     private setIdentity(identity: INotebookIdentity) {
         this.notebookIdentity = identity.resource;
         this.potentialResource = identity.resource;
+        this.notebookType = identity.type;
     }
 
     private async getNotebook(): Promise<INotebook | undefined> {

src/client/datascience/jupyter/jupyterVariables.ts

@@ -15,7 +15,8 @@ import { IFileSystem } from '../../common/platform/types';
 import { IConfigurationService } from '../../common/types';
 import * as localize from '../../common/utils/localize';
 import { EXTENSION_ROOT_DIR } from '../../constants';
-import { Identifiers, Settings } from '../constants';
+import { captureTelemetry } from '../../telemetry';
+import { Identifiers, Settings, Telemetry } from '../constants';
 import {
     ICell,
     IJupyterVariable,
@@ -48,6 +49,7 @@ interface INotebookState {
 export class JupyterVariables implements IJupyterVariables {
     private fetchDataFrameInfoScript?: string;
     private fetchDataFrameRowsScript?: string;
+    private fetchVariableShapeScript?: string;
     private filesLoaded: boolean = false;
     private languageToQueryMap = new Map<string, { query: string; parser: RegExp }>();
     private notebookState = new Map<Uri, INotebookState>();
@@ -58,6 +60,7 @@ export class JupyterVariables implements IJupyterVariables {
     ) {}
 
     // IJupyterVariables implementation
+    @captureTelemetry(Telemetry.VariableExplorerFetchTime, undefined, true)
     public async getVariables(
         notebook: INotebook,
         request: IJupyterVariablesRequest
@@ -102,6 +105,9 @@ export class JupyterVariables implements IJupyterVariables {
         );
         this.fetchDataFrameInfoScript = await this.fileSystem.readFile(file);
 
+        file = path.join(EXTENSION_ROOT_DIR, 'pythonFiles', 'vscode_datascience_helpers', 'getJupyterVariableShape.py');
+        this.fetchVariableShapeScript = await this.fileSystem.readFile(file);
+
         file = path.join(
             EXTENSION_ROOT_DIR,
             'pythonFiles',
@@ -361,7 +367,7 @@ export class JupyterVariables implements IJupyterVariables {
         targetVariable: IJupyterVariable,
         notebook: INotebook
     ): Promise<IJupyterVariable> {
-        const result = { ...targetVariable };
+        let result = { ...targetVariable };
         if (notebook) {
             const output = await notebook.inspect(targetVariable.name);
 
@@ -411,6 +417,26 @@ export class JupyterVariables implements IJupyterVariables {
             }
         }
 
+        // For a python kernel, we might be able to get a better shape. It seems the 'inspect' request doesn't always return it.
+        // Do this only when necessary as this is a LOT slower than an inspect request. Like 4 or 5 times as slow
+        if (
+            result.type &&
+            result.count &&
+            !result.shape &&
+            notebook.getKernelSpec()?.language === 'python' &&
+            result.supportsDataExplorer &&
+            result.type !== 'list' // List count is good enough
+        ) {
+            const computedShape = await this.runScript<IJupyterVariable>(
+                notebook,
+                result,
+                result,
+                () => this.fetchVariableShapeScript
+            );
+            // Only want shape and count from the request. Other things might have been destroyed
+            result = { ...result, shape: computedShape.shape, count: computedShape.count };
+        }
+
         return result;
     }
 }

src/test/datascience/intellisense.unit.test.ts

@@ -5,6 +5,7 @@ import { expect } from 'chai';
 import * as TypeMoq from 'typemoq';
 import * as uuid from 'uuid/v4';
 
+import { Uri } from 'vscode';
 import { IWorkspaceService } from '../../client/common/application/types';
 import { PythonSettings } from '../../client/common/configSettings';
 import { IFileSystem } from '../../client/common/platform/types';
@@ -40,7 +41,7 @@ df
 `;
 
 // tslint:disable-next-line: max-func-body-length
-suite('Data Science Intellisense Unit Tests', () => {
+suite('DataScience Intellisense Unit Tests', () => {
     let intellisenseProvider: IntellisenseProvider;
     let intellisenseDocument: IntellisenseDocument;
     let interpreterService: TypeMoq.IMock<IInterpreterService>;
@@ -291,6 +292,12 @@ suite('Data Science Intellisense Unit Tests', () => {
 
     function loadAllCells(allCells: ICell[]): Promise<void> {
         cells = allCells;
+        intellisenseProvider.onMessage(InteractiveWindowMessages.NotebookIdentity, {
+            resource: Uri.parse('file:///foo.ipynb'),
+            type: 'native'
+        });
+
+        // Load all cells will actually respond with a notification, NotebookIdentity won't so don't wait for it.
         return sendMessage(InteractiveWindowMessages.LoadAllCellsComplete, { cells });
     }
 

src/test/datascience/testHelpers.tsx

@@ -156,6 +156,16 @@ export function runDoubleTest(
     test(`${name} (native)`, async () => testInnerLoop(name, (ioc) => mountWebView(ioc, 'native'), testFunc, getIOC));
 }
 
+export function runInteractiveTest(
+    name: string,
+    testFunc: (wrapper: ReactWrapper<any, Readonly<{}>, React.Component>) => Promise<void>,
+    getIOC: () => DataScienceIocContainer
+) {
+    // Run the test with just the interactive window
+    test(`${name} (interactive)`, async () =>
+        testInnerLoop(name, (ioc) => mountWebView(ioc, 'interactive'), testFunc, getIOC));
+}
+
 export function mountWebView(
     ioc: DataScienceIocContainer,
     type: 'native' | 'interactive'

src/test/datascience/variableexplorer.functional.test.tsx

@@ -15,7 +15,7 @@ import { CommonActionType } from '../../datascience-ui/interactive-common/redux/
 import { DataScienceIocContainer } from './dataScienceIocContainer';
 import { addCode } from './interactiveWindowTestHelpers';
 import { addCell, createNewEditor } from './nativeEditorTestHelpers';
-import { runDoubleTest, waitForMessage } from './testHelpers';
+import { runDoubleTest, runInteractiveTest, waitForMessage } from './testHelpers';
 
 // tslint:disable: no-var-requires no-require-imports
 const rangeInclusive = require('range-inclusive');
@@ -92,7 +92,7 @@ suite('DataScience Interactive Window variable explorer tests', () => {
         }
     }
 
-    runDoubleTest(
+    runInteractiveTest(
         'Variable explorer - Exclude',
         async (wrapper) => {
             const basicCode: string = `import numpy as np
@@ -158,7 +158,7 @@ value = 'hello world'`;
         }
     );
 
-    runDoubleTest(
+    runInteractiveTest(
         'Variable explorer - Update',
         async (wrapper) => {
             const basicCode: string = `value = 'hello world'`;
@@ -254,7 +254,7 @@ value = 'hello world'`;
     );
 
     // Test our display of basic types. We render 8 rows by default so only 8 values per test
-    runDoubleTest(
+    runInteractiveTest(
         'Variable explorer - Types A',
         async (wrapper) => {
             const basicCode: string = `myList = [1, 2, 3]
@@ -285,7 +285,7 @@ myDict = {'a': 1}`;
                     type: 'dict',
                     size: 54,
                     shape: '',
-                    count: 0,
+                    count: 1,
                     truncated: false
                 },
                 {
@@ -295,7 +295,7 @@ myDict = {'a': 1}`;
                     type: 'list',
                     size: 54,
                     shape: '',
-                    count: 0,
+                    count: 3,
                     truncated: false
                 },
                 // Set can vary between python versions, so just don't both to check the value, just see that we got it
@@ -306,7 +306,7 @@ myDict = {'a': 1}`;
                     type: 'set',
                     size: 54,
                     shape: '',
-                    count: 0,
+                    count: 1,
                     truncated: false
                 }
             ];
@@ -317,7 +317,7 @@ myDict = {'a': 1}`;
         }
     );
 
-    runDoubleTest(
+    runInteractiveTest(
         'Variable explorer - Basic B',
         async (wrapper) => {
             const basicCode: string = `import numpy as np
@@ -366,7 +366,7 @@ myTuple = 1,2,3,4,5,6,7,8,9
                     supportsDataExplorer: true,
                     type: 'DataFrame',
                     size: 54,
-                    shape: '',
+                    shape: '(3, 1)',
                     count: 0,
                     truncated: false
                 },
@@ -400,7 +400,7 @@ Name: 0, dtype: float64`,
                     supportsDataExplorer: true,
                     type: 'Series',
                     size: 54,
-                    shape: '',
+                    shape: '(3,)',
                     count: 0,
                     truncated: false
                 },
@@ -410,7 +410,7 @@ Name: 0, dtype: float64`,
                     supportsDataExplorer: false,
                     type: 'tuple',
                     size: 54,
-                    shape: '',
+                    shape: '9',
                     count: 0,
                     truncated: false
                 },
@@ -420,7 +420,7 @@ Name: 0, dtype: float64`,
                     supportsDataExplorer: true,
                     type: 'ndarray',
                     size: 54,
-                    shape: '',
+                    shape: '(3,)',
                     count: 0,
                     truncated: false
                 }
@@ -450,7 +450,8 @@ Name: 0, dtype: float64`,
         };
     }
 
-    // Test our limits. Create 1050 items.
+    // Test our limits. Create 1050 items. Do this with both to make
+    // sure no perf problems with one or the other and to smoke test the native editor
     runDoubleTest(
         'Variable explorer - A lot of items',
         async (wrapper) => {