Get variable info with background thread

pythonFiles/vscode_datascience_helpers/getVariableInfo/vscodeGetVariablesBackground.py

@@ -0,0 +1,179 @@
+# Query Jupyter server for the info about a dataframe
+from collections import namedtuple
+from importlib.util import find_spec
+import json
+
+maxStringLength = 1000
+collectionTypes = ["list", "tuple", "set"]
+arrayPageSize = 50
+
+
+def truncateString(variable):
+    string = repr(variable)
+    if len(string) > maxStringLength:
+        sizeInfo = "\n\nLength: " + str(len(variable)) if type(variable) == str else ""
+        return string[: maxStringLength - 1] + "..." + sizeInfo
+    else:
+        return string
+
+
+DisplayOptions = namedtuple("DisplayOptions", ["width", "max_columns"])
+
+
+def set_pandas_display_options(display_options=None):
+    if find_spec("pandas") is not None:
+        try:
+            import pandas as _VSCODE_PD
+
+            original_display = DisplayOptions(
+                width=_VSCODE_PD.options.display.width,
+                max_columns=_VSCODE_PD.options.display.max_columns,
+            )
+
+            if display_options:
+                _VSCODE_PD.options.display.max_columns = display_options.max_columns
+                _VSCODE_PD.options.display.width = display_options.width
+            else:
+                _VSCODE_PD.options.display.max_columns = 100
+                _VSCODE_PD.options.display.width = 1000
+
+            return original_display
+        except ImportError:
+            pass
+        finally:
+            del _VSCODE_PD
+
+
+def getValue(variable):
+    original_display = None
+    if type(variable).__name__ == "DataFrame" and find_spec("pandas") is not None:
+        original_display = set_pandas_display_options()
+
+    try:
+        return truncateString(variable=variable)
+    finally:
+        if original_display:
+            set_pandas_display_options(original_display)
+
+
+def getPropertyNames(variable):
+    props = []
+    for prop in dir(variable):
+        if not prop.startswith("_"):
+            props.append(prop)
+    return props
+
+
+def getFullType(varType):
+    module = ""
+    if hasattr(varType, "__module__") and varType.__module__ != "builtins":
+        module = varType.__module__ + "."
+    if hasattr(varType, "__qualname__"):
+        return module + varType.__qualname__
+    elif hasattr(varType, "__name__"):
+        return module + varType.__name__
+
+
+def getVariableDescription(variable):
+    result = {}
+
+    varType = type(variable)
+    result["type"] = getFullType(varType)
+    if hasattr(varType, "__mro__"):
+        result["interfaces"] = [getFullType(t) for t in varType.__mro__]
+
+    if hasattr(variable, "__len__") and result["type"] in collectionTypes:
+        result["count"] = len(variable)
+
+    result["hasNamedChildren"] = hasattr(variable, "__dict__") or type(variable) == dict
+
+    result["value"] = getValue(variable)
+    return result
+
+
+def getChildProperty(root, propertyChain):
+    try:
+        variable = root
+        for property in propertyChain:
+            if type(property) == int:
+                if hasattr(variable, "__getitem__"):
+                    variable = variable[property]
+                elif type(variable) == set:
+                    variable = list(variable)[property]
+                else:
+                    return None
+            elif hasattr(variable, property):
+                variable = getattr(variable, property)
+            elif type(variable) == dict and property in variable:
+                variable = variable[property]
+            else:
+                return None
+    except Exception:
+        return None
+
+    return variable
+
+
+### Get info on variables at the root level
+def _VSCODE_getVariableDescriptions(varNames):
+    variables = [
+        {
+            "name": varName,
+            **getVariableDescription(globals()[varName]),
+            "root": varName,
+            "propertyChain": [],
+            "language": "python",
+        }
+        for varName in varNames
+        if varName in globals()
+    ]
+
+    return json.dumps(variables)
+
+
+### Get info on children of a variable reached through the given property chain
+def _VSCODE_getAllChildrenDescriptions(rootVarName, propertyChain, startIndex):
+    root = globals()[rootVarName]
+    if root is None:
+        return []
+
+    parent = root
+    if len(propertyChain) > 0:
+        parent = getChildProperty(root, propertyChain)
+
+    children = []
+    parentInfo = getVariableDescription(parent)
+    if "count" in parentInfo:
+        if parentInfo["count"] > 0:
+            lastItem = min(parentInfo["count"], startIndex + arrayPageSize)
+            indexRange = range(startIndex, lastItem)
+            children = [
+                {
+                    **getVariableDescription(getChildProperty(parent, [i])),
+                    "name": str(i),
+                    "root": rootVarName,
+                    "propertyChain": propertyChain + [i],
+                    "language": "python",
+                }
+                for i in indexRange
+            ]
+    elif parentInfo["hasNamedChildren"]:
+        childrenNames = []
+        if hasattr(parent, "__dict__"):
+            childrenNames = getPropertyNames(parent)
+        elif type(parent) == dict:
+            childrenNames = list(parent.keys())
+
+        children = []
+        for prop in childrenNames:
+            child_property = getChildProperty(parent, [prop])
+            if child_property is not None and type(child_property).__name__ != "method":
+                child = {
+                    **getVariableDescription(child_property),
+                    "name": prop,
+                    "root": rootVarName,
+                    "propertyChain": propertyChain + [prop],
+                }
+                children.append(child)
+
+    return json.dumps(children)

src/kernels/jupyter/types.ts

@@ -16,7 +16,8 @@ import {
     IJupyterConnection,
     GetServerOptions,
     LiveRemoteKernelConnectionMetadata,
-    RemoteKernelConnectionMetadata
+    RemoteKernelConnectionMetadata,
+    IKernel
 } from '../types';
 import { ClassType } from '../../platform/ioc/types';
 import { ContributedKernelFinderKind, IContributedKernelFinder } from '../internalTypes';
@@ -257,3 +258,12 @@ export interface IJupyterServerProviderRegistry {
         serverProvider: JupyterServerProvider
     ): JupyterServerCollection;
 }
+
+export const IBackgroundThreadService = Symbol('IBackgroundThreadService');
+export interface IBackgroundThreadService {
+    execCodeInBackgroundThread<T>(
+        kernel: IKernel,
+        codeWithReturnStatement: string[],
+        token: CancellationToken
+    ): Promise<T | undefined>;
+}

src/kernels/variables/pythonVariableRequester.ts

@@ -14,6 +14,7 @@ import { IKernel } from '../types';
 import { IKernelVariableRequester, IJupyterVariable, IVariableDescription } from './types';
 import { IDataFrameScriptGenerator, IVariableScriptGenerator } from '../../platform/common/types';
 import { SessionDisposedError } from '../../platform/errors/sessionDisposedError';
+import { IBackgroundThreadService } from '../jupyter/types';
 
 type DataFrameSplitFormat = {
     index: (number | string)[];
@@ -75,7 +76,8 @@ async function safeExecuteSilently(
 export class PythonVariablesRequester implements IKernelVariableRequester {
     constructor(
         @inject(IVariableScriptGenerator) private readonly varScriptGenerator: IVariableScriptGenerator,
-        @inject(IDataFrameScriptGenerator) private readonly dfScriptGenerator: IDataFrameScriptGenerator
+        @inject(IDataFrameScriptGenerator) private readonly dfScriptGenerator: IDataFrameScriptGenerator,
+        @inject(IBackgroundThreadService) private readonly backgroundThreadService: IBackgroundThreadService
     ) {}
 
     public async getDataFrameInfo(
@@ -177,35 +179,27 @@ export class PythonVariablesRequester implements IKernelVariableRequester {
         kernel: IKernel,
         parent: IVariableDescription | undefined,
         startIndex: number,
-        token?: CancellationToken
+        token: CancellationToken
     ): Promise<IVariableDescription[]> {
         if (!kernel.session) {
             return [];
         }
 
-        const { code, cleanupCode, initializeCode } =
-            await this.varScriptGenerator.generateCodeToGetAllVariableDescriptions({
-                isDebugging: false,
-                parent,
-                startIndex
-            });
+        const options = parent ? { root: parent.root, propertyChain: parent.propertyChain, startIndex } : undefined;
+        const code = await this.varScriptGenerator.generateCodeToGetAllVariableDescriptions(options);
 
-        const results = await safeExecuteSilently(
+        const content = await this.backgroundThreadService.execCodeInBackgroundThread<IVariableDescription[]>(
             kernel,
-            { code, cleanupCode, initializeCode },
-            {
-                traceErrors: true,
-                traceErrorsMessage: 'Failure in execute_request when retrieving variables',
-                telemetryName: Telemetry.PythonVariableFetchingCodeFailure
-            }
+            code.split(/\r?\n/),
+            token
         );
 
-        if (kernel.disposed || kernel.disposing || token?.isCancellationRequested) {
+        if (kernel.disposed || kernel.disposing || token?.isCancellationRequested || !content) {
             return [];
         }
 
         try {
-            return this.deserializeJupyterResult(results) as Promise<IVariableDescription[]>;
+            return content;
         } catch (ex) {
             traceError(ex);
             return [];

src/platform/common/types.ts

@@ -296,6 +296,11 @@ type ScriptCode = {
      */
     cleanupCode?: string;
 };
+export type ParentOptions = {
+    root: string;
+    propertyChain: (string | number)[];
+    startIndex: number;
+};
 export interface IVariableScriptGenerator {
     generateCodeToGetVariableInfo(options: { isDebugging: boolean; variableName: string }): Promise<ScriptCode>;
     generateCodeToGetVariableProperties(options: {
@@ -304,11 +309,7 @@ export interface IVariableScriptGenerator {
         stringifiedAttributeNameList: string;
     }): Promise<ScriptCode>;
     generateCodeToGetVariableTypes(options: { isDebugging: boolean }): Promise<ScriptCode>;
-    generateCodeToGetAllVariableDescriptions(options: {
-        isDebugging: boolean;
-        parent: { root: string; propertyChain: (string | number)[] } | undefined;
-        startIndex: number;
-    }): Promise<ScriptCode>;
+    generateCodeToGetAllVariableDescriptions(parentOptions: ParentOptions | undefined): Promise<string>;
     generateCodeToGetVariableValueSummary(options: { variableName: string }): Promise<ScriptCode>;
 }
 export const IDataFrameScriptGenerator = Symbol('IDataFrameScriptGenerator');

src/platform/interpreter/variableScriptGenerator.ts

@@ -3,7 +3,7 @@
 
 import { inject, injectable } from 'inversify';
 import { IFileSystem } from '../common/platform/types';
-import { IExtensionContext, IVariableScriptGenerator } from '../common/types';
+import { IExtensionContext, IVariableScriptGenerator, ParentOptions } from '../common/types';
 import { joinPath } from '../vscode-path/resources';
 import dedent from 'dedent';
 
@@ -21,6 +21,7 @@ const cleanupCode = dedent`
 @injectable()
 export class VariableScriptGenerator implements IVariableScriptGenerator {
     static contentsOfScript: string | undefined;
+    static contentsOfBgScript: string | undefined;
     constructor(
         @inject(IFileSystem) private readonly fs: IFileSystem,
         @inject(IExtensionContext) private readonly context: IExtensionContext
@@ -64,38 +65,19 @@ export class VariableScriptGenerator implements IVariableScriptGenerator {
             };
         }
     }
-    async generateCodeToGetAllVariableDescriptions(options: {
-        isDebugging: boolean;
-        parent: { root: string; propertyChain: (string | number)[] } | undefined;
-        startIndex: number;
-    }) {
-        const scriptCode = await this.getContentsOfScript();
-        const isDebugging = options.isDebugging ? 'True' : 'False';
-        const initializeCode = options.parent ? scriptCode : `${scriptCode}\n\n_VSCODE_rwho_ls = %who_ls\n`;
-        const cleanupWhoLsCode = dedent`
-        try:
-            del _VSCODE_rwho_ls
-        except:
-            pass
-        `;
 
-        const code = options.parent
-            ? `${VariableFunc}("AllChildrenDescriptions", ${isDebugging}, "${options.parent.root}", ${JSON.stringify(
-                  options.parent.propertyChain
-              )}, ${options.startIndex})`
-            : `${VariableFunc}("AllVariableDescriptions", ${isDebugging}, _VSCODE_rwho_ls)`;
-        if (options.isDebugging) {
-            return {
-                initializeCode,
-                code,
-                cleanupCode: options.parent ? cleanupCode : `${cleanupCode}\n${cleanupWhoLsCode}`
-            };
+    async generateCodeToGetAllVariableDescriptions(parentOptions: ParentOptions | undefined) {
+        let scriptCode = await this.getContentsOfBgScript();
+        if (parentOptions) {
+            scriptCode =
+                scriptCode +
+                `\n\nreturn _VSCODE_getAllChildrenDescriptions(\'${parentOptions.root}\', ${JSON.stringify(parentOptions.propertyChain)}, ${parentOptions.startIndex})`;
         } else {
-            return {
-                code: `${initializeCode}\n\n${code}\n\n${cleanupCode}`
-            };
+            scriptCode = scriptCode + '\n\nvariables= %who_ls\nreturn _VSCODE_getVariableDescriptions(variables)';
         }
+        return scriptCode;
     }
+
     async generateCodeToGetVariableTypes(options: { isDebugging: boolean }) {
         const scriptCode = await this.getContentsOfScript();
         const initializeCode = `${scriptCode}\n\n_VSCODE_rwho_ls = %who_ls\n`;
@@ -146,4 +128,20 @@ export class VariableScriptGenerator implements IVariableScriptGenerator {
         VariableScriptGenerator.contentsOfScript = contents;
         return contents;
     }
+
+    private async getContentsOfBgScript() {
+        if (VariableScriptGenerator.contentsOfBgScript) {
+            return VariableScriptGenerator.contentsOfBgScript;
+        }
+        const scriptPath = joinPath(
+            this.context.extensionUri,
+            'pythonFiles',
+            'vscode_datascience_helpers',
+            'getVariableInfo',
+            'vscodeGetVariablesBackground.py'
+        );
+        const contents = await this.fs.readFile(scriptPath);
+        VariableScriptGenerator.contentsOfScript = contents;
+        return contents;
+    }
 }

src/standalone/api/kernels/backgroundExecution.ts

@@ -10,6 +10,19 @@ import { raceCancellation } from '../../../platform/common/cancellation';
 import { getNotebookCellOutputMetadata } from '../../../kernels/execution/helpers';
 import { unTrackDisplayDataForExtension } from '../../../kernels/execution/extensionDisplayDataTracker';
 import { traceWarning } from '../../../platform/logging';
+import { IBackgroundThreadService } from '../../../kernels/jupyter/types';
+import { injectable } from 'inversify';
+
+@injectable()
+export class BackgroundThreadService implements IBackgroundThreadService {
+    execCodeInBackgroundThread<T>(
+        kernel: IKernel,
+        codeWithReturnStatement: string[],
+        token: CancellationToken
+    ): Promise<T | undefined> {
+        return execCodeInBackgroundThread(kernel, codeWithReturnStatement, token);
+    }
+}
 
 export const executionCounters = new WeakMap<IKernel, number>();
 export async function execCodeInBackgroundThread<T>(