Add API to get notebook URI connected to a text document URI (#10865)

Author: debonte

src/platform/api/types.ts

@@ -125,7 +125,21 @@ export type PythonApi = {
         interpreter?: PythonEnvironment_PythonApi
     ): Promise<string[] | undefined>;
 
+    /**
+     * Call to provide a function that the Python extension can call to request the Python
+     * path to use for a particular notebook.
+     * @param func : The function that Python should call when requesting the Python path.
+     */
     registerJupyterPythonPathFunction(func: (uri: Uri) => Promise<string | undefined>): void;
+
+    /**
+     * Call to provide a function that the Python extension can call to request the notebook
+     * document URI related to a particular text document URI, or undefined if there is no
+     * associated notebook.
+     * @param func : The function that Python should call when requesting the notebook URI.
+     */
+    registerGetNotebookUriForTextDocumentUriFunction(func: (textDocumentUri: Uri) => Uri | undefined): void;
+
     /**
      * This API will re-trigger environment discovery. Extensions can wait on the returned
      * promise to get the updated interpreters list. If there is a refresh already going on

src/standalone/intellisense/notebookPythonPathService.node.ts

@@ -47,6 +47,11 @@ export class NotebookPythonPathService implements IExtensionSingleActivationServ
             if (api.registerJupyterPythonPathFunction !== undefined) {
                 api.registerJupyterPythonPathFunction((uri) => this._jupyterPythonPathFunction(uri));
             }
+            if (api.registerGetNotebookUriForTextDocumentUriFunction !== undefined) {
+                api.registerGetNotebookUriForTextDocumentUriFunction((uri) =>
+                    this._getNotebookUriForTextDocumentUri(uri)
+                );
+            }
         });
     }
 
@@ -138,4 +143,14 @@ export class NotebookPythonPathService implements IExtensionSingleActivationServ
 
         return pythonPath;
     }
+
+    private _getNotebookUriForTextDocumentUri(textDocumentUri: Uri): Uri | undefined {
+        if (textDocumentUri.scheme !== 'vscode-interactive-input') {
+            return undefined;
+        }
+
+        const notebookPath = `${textDocumentUri.fsPath.replace('\\InteractiveInput-', 'Interactive-')}.interactive`;
+        const notebookUri = textDocumentUri.with({ scheme: 'vscode-interactive', path: notebookPath });
+        return notebookUri;
+    }
 }