From f7e073fa0556efd694e78fae81b4fb773cd44486 Mon Sep 17 00:00:00 2001
From: Don Jayamanne <don.jayamanne@yahoo.com>
Date: Wed, 9 Sep 2020 15:03:01 -0700
Subject: [PATCH] Find python based on hash in kernelspec

---
 .../jupyter/kernels/kernelSelector.ts         | 26 +++++++++++++++++++
 .../datascience/notebookStorage/baseModel.ts  | 15 ++++++++++-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/client/datascience/jupyter/kernels/kernelSelector.ts b/src/client/datascience/jupyter/kernels/kernelSelector.ts
index d2623d2df616..df8106566a92 100644
--- a/src/client/datascience/jupyter/kernels/kernelSelector.ts
+++ b/src/client/datascience/jupyter/kernels/kernelSelector.ts
@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 import type { nbformat } from '@jupyterlab/coreutils';
 import type { Kernel } from '@jupyterlab/services';
+import { sha256 } from 'hash.js';
 import { inject, injectable } from 'inversify';
 // tslint:disable-next-line: no-require-imports
 import cloneDeep = require('lodash/cloneDeep');
@@ -18,6 +19,7 @@ import { PythonEnvironment } from '../../../pythonEnvironments/info';
 import { IEventNamePropertyMapping, sendTelemetryEvent } from '../../../telemetry';
 import { Commands, KnownNotebookLanguages, Settings, Telemetry } from '../../constants';
 import { IKernelFinder } from '../../kernel-launcher/types';
+import { getInterpreterInfoStoredInMetadata } from '../../notebookStorage/baseModel';
 import { reportAction } from '../../progress/decorator';
 import { ReportableAction } from '../../progress/types';
 import {
@@ -485,6 +487,17 @@ export class KernelSelector implements IKernelSelectionUsage {
             }
         }
     }
+    private async findInterpreterStoredInNotebookMetadata(
+        resource: Resource,
+        notebookMetadata?: nbformat.INotebookMetadata
+    ): Promise<PythonEnvironment | undefined> {
+        const info = getInterpreterInfoStoredInMetadata(notebookMetadata);
+        if (!info) {
+            return;
+        }
+        const interpreters = await this.interpreterService.getInterpreters(resource);
+        return interpreters.find((item) => sha256().update(item.path).digest('hex') === info.hash);
+    }
 
     // Get our kernelspec and interpreter for a local raw connection
     private async getKernelForLocalRawConnection(
@@ -493,6 +506,19 @@ export class KernelSelector implements IKernelSelectionUsage {
         cancelToken?: CancellationToken,
         ignoreDependencyCheck?: boolean
     ): Promise<KernelSpecConnectionMetadata | PythonKernelConnectionMetadata | undefined> {
+        // If user had selected an interpreter (raw kernel), then that interpreter would be stored in the kernelspec metadata.
+        // Find this matching interpreter & start that using raw kernel.
+        const interpreterStoredInKernelSpec = await this.findInterpreterStoredInNotebookMetadata(
+            resource,
+            notebookMetadata
+        );
+        if (interpreterStoredInKernelSpec) {
+            return {
+                kind: 'startUsingPythonInterpreter',
+                interpreter: interpreterStoredInKernelSpec
+            };
+        }
+
         // First use our kernel finder to locate a kernelspec on disk
         const kernelSpec = await this.kernelFinder.findKernelSpec(
             resource,
diff --git a/src/client/datascience/notebookStorage/baseModel.ts b/src/client/datascience/notebookStorage/baseModel.ts
index 2a28e231352f..dc5a5322f143 100644
--- a/src/client/datascience/notebookStorage/baseModel.ts
+++ b/src/client/datascience/notebookStorage/baseModel.ts
@@ -23,6 +23,19 @@ type KernelIdListEntry = {
     kernelId: string | undefined;
 };
 
+export function getInterpreterInfoStoredInMetadata(
+    metadata?: nbformat.INotebookMetadata
+): { displayName: string; hash: string } | undefined {
+    if (!metadata || !metadata.kernelspec || !metadata.kernelspec.name) {
+        return;
+    }
+    // See `updateNotebookMetadata` to determine how & where exactly interpreter hash is stored.
+    // tslint:disable-next-line: no-any
+    const kernelSpecMetadata: undefined | any = metadata.kernelspec.metadata as any;
+    const interpreterHash = kernelSpecMetadata?.interpreter?.hash;
+    return interpreterHash ? { displayName: metadata.kernelspec.name, hash: interpreterHash } : undefined;
+}
+
 // tslint:disable-next-line: cyclomatic-complexity
 export function updateNotebookMetadata(
     metadata?: nbformat.INotebookMetadata,
@@ -80,7 +93,7 @@ export function updateNotebookMetadata(
         }
     } else if (kernelConnection?.kind === 'startUsingPythonInterpreter') {
         // Store interpreter name, we expect the kernel finder will find the corresponding interpreter based on this name.
-        const name = kernelConnection.interpreter.displayName || kernelConnection.interpreter.path;
+        const name = kernelConnection.interpreter.displayName || '';
         if (metadata.kernelspec?.name !== name || metadata.kernelspec?.display_name !== name) {
             changed = true;
             metadata.kernelspec = {