Move getAssociatedNotebookDocument into kernel helpers (#10349)

TELEMETRY.md

@@ -2805,7 +2805,7 @@ No properties for event
                 : this.start(new DisplayOptions(false)));
             sendKernelTelemetryEvent(this.resourceUri, Telemetry.NotebookRestart, stopWatch.elapsedTime);
         } catch (ex) {
-            traceError(`Restart failed ${getDisplayPath(this.id)}`, ex);
+            traceError(`Restart failed ${getDisplayPath(this.uri)}`, ex);
             this._ignoreJupyterSessionDisposedErrors = true;
 ```
 

src/interactive-window/debugger/interactiveWindowDebugger.node.ts

@@ -16,8 +16,7 @@ import { IKernel, isLocalConnection } from '../../kernels/types';
 import { IInteractiveWindowDebugger } from '../types';
 import { IFileGeneratedCodes } from '../editor-integration/types';
 import { IJupyterDebugService } from '../../kernels/debugger/types';
-import { getAssociatedNotebookDocument } from '../../notebooks/controllers/kernelSelector';
-import { executeSilently } from '../../kernels/helpers';
+import { executeSilently, getAssociatedNotebookDocument } from '../../kernels/helpers';
 import { buildSourceMap } from './helper';
 
 @injectable()

src/interactive-window/editor-integration/codeLensFactory.ts

@@ -25,7 +25,7 @@ import { InteractiveWindowView } from '../../notebooks/constants';
 import { CodeLensCommands, Commands } from '../../platform/common/constants';
 import { generateCellRangesFromDocument } from './cellFactory';
 import { ICodeLensFactory, IGeneratedCode, IGeneratedCodeStorageFactory } from './types';
-import { getAssociatedNotebookDocument } from '../../notebooks/controllers/kernelSelector';
+import { getAssociatedNotebookDocument } from '../../kernels/helpers';
 
 type CodeLensCacheData = {
     cachedDocumentVersion: number | undefined;

src/interactive-window/generatedCodeStoreManager.ts

@@ -3,9 +3,9 @@
 
 import { inject, injectable } from 'inversify';
 import { NotebookDocument } from 'vscode';
+import { getAssociatedNotebookDocument } from '../kernels/helpers';
 import { IKernel, IKernelProvider } from '../kernels/types';
 import { InteractiveWindowView } from '../notebooks/constants';
-import { getAssociatedNotebookDocument } from '../notebooks/controllers/kernelSelector';
 import { INotebookControllerManager } from '../notebooks/types';
 import { IExtensionSyncActivationService } from '../platform/activation/types';
 import { disposeAllDisposables } from '../platform/common/helpers';

src/kernels/helpers.ts

@@ -16,9 +16,10 @@ import {
     LiveRemoteKernelConnectionMetadata,
     PythonKernelConnectionMetadata,
     IJupyterKernelSpec,
-    IJupyterSession
+    IJupyterSession,
+    IKernel
 } from './types';
-import { Uri } from 'vscode';
+import { Uri, workspace } from 'vscode';
 import { IWorkspaceService } from '../platform/common/application/types';
 import { isCI, PYTHON_LANGUAGE, Telemetry } from '../platform/common/constants';
 import { traceError, traceInfo, traceInfoIfCI, traceWarning } from '../platform/logging';
@@ -1574,3 +1575,11 @@ export function deserializeKernelConnection(kernelConnection: any): KernelConnec
     }
     return kernelConnection;
 }
+
+export function getAssociatedNotebookDocument(kernel: IKernel | undefined) {
+    if (!kernel) {
+        return;
+    }
+
+    return workspace.notebookDocuments.find((nb) => nb.uri.toString() === kernel.uri.toString());
+}

src/kernels/ipywidgets-message-coordination/ipyWidgetMessageDispatcher.ts

@@ -18,7 +18,7 @@ import { Identifiers, Telemetry } from '../../webviews/webview-side/common/const
 import { IKernel, IKernelProvider, KernelSocketInformation } from '../types';
 import { WIDGET_MIMETYPE } from './constants';
 import { IIPyWidgetMessageDispatcher, IPyWidgetMessage } from './types';
-import { getAssociatedNotebookDocument } from '../../notebooks/controllers/kernelSelector';
+import { getAssociatedNotebookDocument } from '../helpers';
 
 type PendingMessage = {
     resultPromise: Deferred<void>;

src/kernels/ipywidgets-message-coordination/ipyWidgetMessageDispatcherFactory.ts

@@ -5,9 +5,9 @@
 
 import { inject, injectable } from 'inversify';
 import { Event, EventEmitter, NotebookDocument } from 'vscode';
-import { getAssociatedNotebookDocument } from '../../notebooks/controllers/kernelSelector';
 import { IDisposable, IDisposableRegistry } from '../../platform/common/types';
 import { IPyWidgetMessages } from '../../platform/messageTypes';
+import { getAssociatedNotebookDocument } from '../helpers';
 import { IKernel, IKernelProvider } from '../types';
 import { IPyWidgetMessageDispatcher } from './ipyWidgetMessageDispatcher';
 import { IIPyWidgetMessageDispatcher, IPyWidgetMessage } from './types';

src/kernels/ipywidgets-message-coordination/ipyWidgetScriptSource.ts

@@ -19,8 +19,8 @@ import { Telemetry } from '../../webviews/webview-side/common/constants';
 import { IKernel, IKernelProvider } from '../types';
 import { IPyWidgetScriptSourceProvider } from './ipyWidgetScriptSourceProvider';
 import { ILocalResourceUriConverter, IWidgetScriptSourceProviderFactory, WidgetScriptSource } from './types';
-import { getAssociatedNotebookDocument } from '../../notebooks/controllers/kernelSelector';
 import { ConsoleForegroundColors } from '../../platform/logging/types';
+import { getAssociatedNotebookDocument } from '../helpers';
 
 export class IPyWidgetScriptSource {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any

src/kernels/kernel.base.ts

@@ -133,7 +133,7 @@ export abstract class BaseKernel implements IKernel {
     private disposingPromise?: Promise<void>;
 
     constructor(
-        public readonly id: Uri,
+        public readonly uri: Uri,
         public readonly resourceUri: Resource,
         public readonly kernelConnectionMetadata: Readonly<KernelConnectionMetadata>,
         private readonly notebookProvider: INotebookProvider,
@@ -206,17 +206,17 @@ export abstract class BaseKernel implements IKernel {
         await Promise.all(this.eventHooks.map((h) => h('willInterrupt')));
         trackKernelResourceInformation(this.resourceUri, { interruptKernel: true });
         if (this.restarting) {
-            traceInfo(`Interrupt requested & currently restarting ${getDisplayPath(this.resourceUri || this.id)}`);
+            traceInfo(`Interrupt requested & currently restarting ${getDisplayPath(this.resourceUri || this.uri)}`);
             await this.restarting.promise;
         }
-        traceInfo(`Interrupt requested ${getDisplayPath(this.resourceUri || this.id)}`);
+        traceInfo(`Interrupt requested ${getDisplayPath(this.resourceUri || this.uri)}`);
         this.startCancellation.cancel();
         const interruptResultPromise = this.kernelExecution.interrupt(this._jupyterSessionPromise);
 
         const status = this.statusProvider.set(DataScience.interruptKernelStatus());
         let result: InterruptResult | undefined;
         try {
-            traceInfo(`Interrupt requested & sent for ${getDisplayPath(this.id)} in notebookEditor.`);
+            traceInfo(`Interrupt requested & sent for ${getDisplayPath(this.uri)} in notebookEditor.`);
             result = await interruptResultPromise;
             if (result === InterruptResult.TimedOut) {
                 const message = DataScience.restartKernelAfterInterruptMessage();
@@ -232,7 +232,7 @@ export abstract class BaseKernel implements IKernel {
         }
     }
     public async dispose(): Promise<void> {
-        traceInfo(`Dispose Kernel '${getDisplayPath(this.id)}' associated with '${getDisplayPath(this.resourceUri)}'`);
+        traceInfo(`Dispose Kernel '${getDisplayPath(this.uri)}' associated with '${getDisplayPath(this.resourceUri)}'`);
         this._disposing = true;
         if (this.disposingPromise) {
             return this.disposingPromise;
@@ -271,7 +271,7 @@ export abstract class BaseKernel implements IKernel {
             return this.restarting.promise;
         }
         await Promise.all(this.eventHooks.map((h) => h('willRestart')));
-        traceInfo(`Restart requested ${this.id}`);
+        traceInfo(`Restart requested ${this.uri}`);
         this.startCancellation.cancel();
         // Set our status
         const status = this.statusProvider.set(DataScience.restartingKernelStatus().format(''));
@@ -290,7 +290,7 @@ export abstract class BaseKernel implements IKernel {
                 : this.start(new DisplayOptions(false)));
             sendKernelTelemetryEvent(this.resourceUri, Telemetry.NotebookRestart, stopWatch.elapsedTime);
         } catch (ex) {
-            traceError(`Restart failed ${getDisplayPath(this.id)}`, ex);
+            traceError(`Restart failed ${getDisplayPath(this.uri)}`, ex);
             this._ignoreJupyterSessionDisposedErrors = true;
             // If restart fails, kill the associated session.
             const session = this._session;
@@ -377,7 +377,7 @@ export abstract class BaseKernel implements IKernel {
             }
             this._jupyterSessionPromise = this.createJupyterSession(new StopWatch()).catch((ex) => {
                 traceInfoIfCI(
-                    `Failed to create Jupyter Session in Kernel.startNotebook for ${getDisplayPath(this.id)}`
+                    `Failed to create Jupyter Session in Kernel.startNotebook for ${getDisplayPath(this.uri)}`
                 );
                 // If we fail also clear the promise.
                 this.startCancellation.cancel();
@@ -404,7 +404,7 @@ export abstract class BaseKernel implements IKernel {
             traceInfo(
                 `Starting Jupyter Session id = '${this.kernelConnectionMetadata.kind}:${
                     this.kernelConnectionMetadata.id
-                }'${pythonInfo} for '${getDisplayPath(this.id)}' (disableUI=${this.startupUI.disableUI})`
+                }'${pythonInfo} for '${getDisplayPath(this.uri)}' (disableUI=${this.startupUI.disableUI})`
             );
             this.createProgressIndicator(disposables);
             this.isKernelDead = false;
@@ -483,7 +483,7 @@ export abstract class BaseKernel implements IKernel {
     protected abstract sendTelemetryForPythonKernelExecutable(): Promise<void>;
 
     protected async initializeAfterStart(session: IJupyterSession | undefined) {
-        traceVerbose(`Started running kernel initialization for ${getDisplayPath(this.id)}`);
+        traceVerbose(`Started running kernel initialization for ${getDisplayPath(this.uri)}`);
         if (!session) {
             traceVerbose('Not running kernel initialization');
             return;
@@ -494,14 +494,14 @@ export abstract class BaseKernel implements IKernel {
             session.onDidDispose(() => {
                 traceInfoIfCI(
                     `Kernel got disposed as a result of session.onDisposed ${getDisplayPath(
-                        this.resourceUri || this.id
+                        this.resourceUri || this.uri
                     )}`
                 );
                 // Ignore when session is disposed as a result of failed restarts.
                 if (!this._ignoreJupyterSessionDisposedErrors) {
                     traceInfo(
                         `Kernel got disposed as a result of session.onDisposed ${getDisplayPath(
-                            this.resourceUri || this.id
+                            this.resourceUri || this.uri
                         )} & _ignoreJupyterSessionDisposedErrors = false.`
                     );
                     const isActiveSessionDead = this._session === session;
@@ -660,11 +660,11 @@ export abstract class BaseKernel implements IKernel {
         const settings = this.configService.getSettings(this.resourceUri);
         if (settings && settings.themeMatplotlibPlots) {
             // We're theming matplotlibs, so we have to setup our default state.
-            traceInfoIfCI(`Initialize config for plots for ${getDisplayPath(this.resourceUri || this.id)}`);
+            traceInfoIfCI(`Initialize config for plots for ${getDisplayPath(this.resourceUri || this.uri)}`);
 
             const matplotInit = CodeSnippets.MatplotLibInit;
 
-            traceInfo(`Initialize matplotlib for ${getDisplayPath(this.resourceUri || this.id)}`);
+            traceInfo(`Initialize matplotlib for ${getDisplayPath(this.resourceUri || this.uri)}`);
             // Force matplotlib to inline and save the default style. We'll use this later if we
             // get a request to update style
             results.push(...matplotInit.splitLines({ trim: false }));

src/kernels/kernel.node.ts

@@ -12,7 +12,7 @@ import { InteractiveWindowView } from '../notebooks/constants';
 import { calculateWorkingDirectory } from '../platform/common/utils.node';
 import { CodeSnippets } from '../webviews/webview-side/common/constants';
 import { CellOutputDisplayIdTracker } from '../notebooks/execution/cellDisplayIdTracker';
-import { isLocalHostConnection, isPythonKernelConnection } from './helpers';
+import { getAssociatedNotebookDocument, isLocalHostConnection, isPythonKernelConnection } from './helpers';
 import { expandWorkingDir } from './jupyter/jupyterUtils';
 import {
     INotebookProvider,
@@ -23,13 +23,12 @@ import {
 } from './types';
 import { AddRunCellHook } from '../platform/common/scriptConstants';
 import { IStatusProvider } from '../platform/progress/types';
-import { getAssociatedNotebookDocument } from '../notebooks/controllers/kernelSelector';
 import { sendTelemetryForPythonKernelExecutable } from './helpers.node';
 import { BaseKernel } from './kernel.base';
 
 export class Kernel extends BaseKernel {
     constructor(
-        id: Uri,
+        uri: Uri,
         resourceUri: Resource,
         kernelConnectionMetadata: Readonly<KernelConnectionMetadata>,
         notebookProvider: INotebookProvider,
@@ -48,7 +47,7 @@ export class Kernel extends BaseKernel {
         formatters: ITracebackFormatter[]
     ) {
         super(
-            id,
+            uri,
             resourceUri,
             kernelConnectionMetadata,
             notebookProvider,

src/kernels/kernel.web.ts

@@ -10,7 +10,7 @@ import { BaseKernel } from './kernel.base';
 import { CellOutputDisplayIdTracker } from '../notebooks/execution/cellDisplayIdTracker';
 import { IStatusProvider } from '../platform/progress/types';
 import { InteractiveWindowView } from '../notebooks/constants';
-import { getAssociatedNotebookDocument } from '../notebooks/controllers/kernelSelector';
+import { getAssociatedNotebookDocument } from './helpers';
 const addRunCellHook = require('../../pythonFiles/vscode_datascience_helpers/kernel/addRunCellHook.py');
 
 /**

src/kernels/kernelCommandListener.ts

@@ -16,10 +16,10 @@ import { Commands, Telemetry } from '../webviews/webview-side/common/constants';
 import { IKernel, IKernelProvider } from './types';
 import { IInteractiveWindowProvider } from '../interactive-window/types';
 import { IDataScienceErrorHandler } from '../platform/errors/types';
-import { getAssociatedNotebookDocument } from '../notebooks/controllers/kernelSelector';
 import { DisplayOptions } from './displayOptions';
 import { KernelConnector } from './kernelConnector';
 import { getDisplayPath } from '../platform/common/platform/fs-paths';
+import { getAssociatedNotebookDocument } from './helpers';
 
 @injectable()
 export class KernelCommandListener implements IDataScienceCommandListener {

src/kernels/types.ts

@@ -123,7 +123,7 @@ export function isLocalConnection(
 }
 
 export interface IKernel extends IAsyncDisposable {
-    readonly id: Uri;
+    readonly uri: Uri;
     /**
      * In the case of Notebooks, this is the same as the Notebook Uri.
      * But in the case of Interactive Window, this is the Uri of the file (such as the Python file).

src/kernels/variables/debuggerVariables.ts

@@ -7,7 +7,6 @@ import * as uriPath from '../../platform/vscode-path/resources';
 
 import { DebugAdapterTracker, Disposable, Event, EventEmitter } from 'vscode';
 import { DebugProtocol } from 'vscode-debugprotocol';
-import { getAssociatedNotebookDocument } from '../../notebooks/controllers/kernelSelector';
 import { IDebugService, IVSCodeNotebook } from '../../platform/common/application/types';
 import { DataFrameLoading, GetVariableInfo } from '../../platform/common/scriptConstants';
 import { traceError, traceVerbose } from '../../platform/logging';
@@ -27,6 +26,7 @@ import {
 import { convertDebugProtocolVariableToIJupyterVariable, DataViewableTypes } from './helpers';
 import { IFileSystem } from '../../platform/common/platform/types';
 import { noop } from '../../platform/common/utils/misc';
+import { getAssociatedNotebookDocument } from '../helpers';
 
 const KnownExcludedVariables = new Set<string>(['In', 'Out', 'exit', 'quit']);
 const MaximumRowChunkSizeForDebugger = 100;

src/kernels/variables/kernelVariables.ts

@@ -79,7 +79,7 @@ export class KernelVariables implements IJupyterVariables {
         token?: CancellationToken
     ): Promise<IJupyterVariable | undefined> {
         // See if in the cache
-        const cache = this.cachedVariables.get(kernel.id.toString());
+        const cache = this.cachedVariables.get(kernel.uri.toString());
         if (cache) {
             let match = cache.variables.find((v) => v.name === name);
             if (match && !match.value) {
@@ -170,7 +170,7 @@ export class KernelVariables implements IJupyterVariables {
         request: IJupyterVariablesRequest
     ): Promise<IJupyterVariablesResponse> {
         // See if we already have the name list
-        let list = this.cachedVariables.get(kernel.id.toString());
+        let list = this.cachedVariables.get(kernel.uri.toString());
         if (!list || list.currentExecutionCount !== request.executionCount) {
             // Refetch the list of names from the notebook. They might have changed.
             list = {
@@ -241,7 +241,7 @@ export class KernelVariables implements IJupyterVariables {
             }
 
             // Save in our cache
-            this.cachedVariables.set(kernel.id.toString(), list);
+            this.cachedVariables.set(kernel.uri.toString(), list);
 
             // Update total count (exclusions will change this as types are computed)
             result.totalCount = list.variables.length;

src/kernels/variables/pythonVariableRequester.ts

@@ -8,10 +8,9 @@ import { DataScience } from '../../platform/common/utils/localize';
 import { stripAnsi } from '../../platform/common/utils/regexp';
 import { JupyterDataRateLimitError } from '../../platform/errors/jupyterDataRateLimitError';
 import { Telemetry } from '../../webviews/webview-side/common/constants';
-import { executeSilently } from '../helpers';
+import { executeSilently, getAssociatedNotebookDocument } from '../helpers';
 import { IKernel } from '../types';
 import { IKernelVariableRequester, IJupyterVariable } from './types';
-import { getAssociatedNotebookDocument } from '../../notebooks/controllers/kernelSelector';
 import { DataFrameLoading, GetVariableInfo } from '../../platform/common/scriptConstants';
 import { IExtensionContext } from '../../platform/common/types';
 
@@ -68,7 +67,7 @@ export class PythonVariablesRequester implements IKernelVariableRequester {
               )
             : [];
 
-        const fileName = getAssociatedNotebookDocument(kernel)?.uri || kernel.resourceUri || kernel.id;
+        const fileName = getAssociatedNotebookDocument(kernel)?.uri || kernel.resourceUri || kernel.uri;
 
         // Combine with the original result (the call only returns the new fields)
         return {

src/notebooks/controllers/kernelSelector.ts

@@ -4,8 +4,6 @@
 import { ICommandManager } from '../../platform/common/application/types';
 import { traceError } from '../../platform/logging';
 import { Resource } from '../../platform/common/types';
-import { IKernel } from '../../kernels/types';
-import { workspace } from 'vscode';
 import { INotebookEditorProvider } from '../types';
 
 /**
@@ -25,11 +23,3 @@ export async function selectKernel(
     traceError(`Unable to select kernel as the Notebook document could not be identified`);
     return false;
 }
-
-export function getAssociatedNotebookDocument(kernel: IKernel | undefined) {
-    if (!kernel) {
-        return;
-    }
-
-    return workspace.notebookDocuments.find((nb) => nb.uri.toString() === kernel.id.toString());
-}

src/notebooks/execution/kernelExecution.ts

@@ -26,8 +26,8 @@ import {
 } from '../../kernels/types';
 import { traceCellMessage } from '../helpers';
 import { getDisplayPath } from '../../platform/common/platform/fs-paths';
-import { getAssociatedNotebookDocument } from '../controllers/kernelSelector';
 import { CellExecutionMessageHandlerService } from './cellExecutionMessageHandlerService';
+import { getAssociatedNotebookDocument } from '../../kernels/helpers';
 
 /**
  * Separate class that deals just with kernel execution.