feat: VS Code - Generate requirements.txt file for server environment (…

…#196)

DH-18163: VS Code - Generate requirements.txt file for server
environment

Should be able to generate a `requirements.txt` file via a right-click
context menu in an active connection:
<img width="517" alt="image"
src="https://github.com/user-attachments/assets/fa534e08-aff0-4586-9f91-3cc51e83cc06"
/>

README.md

@@ -51,16 +51,16 @@ Enterprise servers can be configured via the `"deephaven.enterpriseServers"` set
 ![Enterprise Server Settings](./docs/assets/dhe-settings.gif)
 
 ## Workspace Setup
-It is recommended to configure a virtual python environment within your `VS Code` workspace. See https://code.visualstudio.com/docs/python/python-tutorial#_create-a-virtual-environment for a general overview. To get features like intellisense, you can install `deephaven` pip packages in the `venv`.
+It is recommended to configure a virtual python environment within your `VS Code` workspace. See https://code.visualstudio.com/docs/python/python-tutorial#_create-a-virtual-environment for a general overview. To get features like intellisense for packages that are installed on the Deephaven server, you will need to install the same packages in your local `venv`.
 
-For example here's a minimal `requirements.txt` file that enables intellisense for common Deephaven packages:
-```text
-deephaven-core
-deephaven-plugin-plotly-express
-deephaven-plugin-ui
-```
+A `requirements.txt` file can be generated containing all of the packages installed on the server by:
+1. Connect to a Deephaven server
+1. Right-click on the connection in the CONNECTIONS panel
+1. Click "Generate requirements.txt" action
+
+   ![Generate requirements.txt](docs/assets/generate-requirements-txt.png)
 
-> Note: Python code executed by the extension always runs on the server, while the local environment drives language features in `VS Code` such as intellisense. Therefore, local pip installs will need to target the same versions installed on the server to ensure intellisense features match the apis running on the server. For Community, it is possible for the server to share the same environment as `VS Code`. For Enterprise, they will always be separate. We plan to implement better support for this in the future.
+> Note: Python code executed by the extension always runs on the server, while the local environment drives language features in `VS Code` such as intellisense. For Community, it is possible for the server to share the same environment as `VS Code`. For Enterprise, they will always be separate.
 
 ### Managed Pip Servers (Community only)
 If you want to manage Deephaven servers from within the extension, include `deephaven-server` in the venv pip installation.

package.json

@@ -194,6 +194,11 @@
         "title": "Generate DHE Key Pair",
         "icon": "$(key)"
       },
+      {
+        "command": "vscode-deephaven.generateRequirementsTxt",
+        "title": "Generate requirements.txt",
+        "icon": "$(cloud-download)"
+      },
       {
         "command": "vscode-deephaven.openInBrowser",
         "title": "Deephaven: Open in Browser",
@@ -696,6 +701,10 @@
           "command": "vscode-deephaven.generateDHEKeyPair",
           "when:": "false"
         },
+        {
+          "command": "vscode-deephaven.generateRequirementsTxt",
+          "when": "false"
+        },
         {
           "command": "vscode-deephaven.openVariablePanels",
           "when": "false"
@@ -782,23 +791,27 @@
         },
         {
           "command": "vscode-deephaven.createNewTextDoc",
-          "when": "view == vscode-deephaven.serverConnectionTree && viewItem == isConnection",
+          "when": "view == vscode-deephaven.serverConnectionTree && viewItem == isConnectionConnected",
           "group": "inline@1"
         },
         {
           "command": "vscode-deephaven.disconnectEditor",
-          "when": "view == vscode-deephaven.serverConnectionTree && viewItem != isConnection",
+          "when": "view == vscode-deephaven.serverConnectionTree && viewItem == isUri",
           "group": "inline"
         },
         {
           "command": "vscode-deephaven.disconnectFromServer",
-          "when": "(view == vscode-deephaven.serverTree && (viewItem == isServerRunningConnected || viewItem == isDHEServerRunningConnected)) || (view == vscode-deephaven.serverConnectionTree && viewItem == isConnection)",
+          "when": "(view == vscode-deephaven.serverTree && (viewItem == isServerRunningConnected || viewItem == isDHEServerRunningConnected)) || (view == vscode-deephaven.serverConnectionTree && (viewItem == isConnectionConnected || viewItem == isConnectionConnecting))",
           "group": "inline@2"
         },
         {
           "command": "vscode-deephaven.generateDHEKeyPair",
           "when": "view == vscode-deephaven.serverTree && (viewItem == isDHEServerRunningConnected || viewItem == isDHEServerRunningDisconnected)"
         },
+        {
+          "command": "vscode-deephaven.generateRequirementsTxt",
+          "when": "view == vscode-deephaven.serverConnectionTree && viewItem == isConnectionConnected"
+        },
         {
           "command": "vscode-deephaven.openInBrowser",
           "when": "view == vscode-deephaven.serverTree && (viewItem == isManagedServerConnected || viewItem == isManagedServerDisconnected || viewItem == isServerRunningConnected || viewItem == isServerRunningDisconnected || viewItem == isDHEServerRunningConnected || viewItem == isDHEServerRunningDisconnected)",

src/common/commands.ts

@@ -22,6 +22,7 @@ export const DISCONNECT_EDITOR_CMD = cmd('disconnectEditor');
 export const DISCONNECT_FROM_SERVER_CMD = cmd('disconnectFromServer');
 export const DOWNLOAD_LOGS_CMD = cmd('downloadLogs');
 export const GENERATE_DHE_KEY_PAIR_CMD = cmd('generateDHEKeyPair');
+export const GENERATE_REQUIREMENTS_TXT_CMD = cmd('generateRequirementsTxt');
 export const OPEN_IN_BROWSER_CMD = cmd('openInBrowser');
 export const OPEN_VARIABLE_PANELS_CMD = cmd('openVariablePanels');
 export const REFRESH_SERVER_TREE_CMD = cmd('refreshServerTree');

src/common/constants.ts

@@ -91,7 +91,8 @@ export const VARIABLE_UNICODE_ICONS = {
 /* eslint-enable @typescript-eslint/naming-convention */
 
 export const CONNECTION_TREE_ITEM_CONTEXT = {
-  isConnection: 'isConnection',
+  isConnectionConnected: 'isConnectionConnected',
+  isConnectionConnecting: 'isConnectionConnecting',
   isUri: 'isUri',
 } as const;
 
@@ -120,3 +121,25 @@ export const VSCODE_POST_MSG = {
   loginOptionsResponse: 'vscode-ext.loginOptions',
   sessionDetailsResponse: 'vscode-ext.sessionDetails',
 } as const;
+
+/**
+ * Table to store Python dependency names + versions used to generate a
+ * requirements.txt file
+ */
+export const REQUIREMENTS_TABLE_NAME = '__vscode_requirements';
+export const REQUIREMENTS_TABLE_NAME_COLUMN_NAME = 'Name';
+export const REQUIREMENTS_TABLE_VERSION_COLUMN_NAME = 'Version';
+
+/**
+ * Query installed Python package names + versions and store in a DH Table.
+ */
+export const REQUIREMENTS_QUERY_TXT = `from deephaven import new_table
+from deephaven.column import string_col
+from importlib.metadata import packages_distributions, version
+
+installed = {pkg for pkgs in packages_distributions().values() for pkg in pkgs}
+
+${REQUIREMENTS_TABLE_NAME} = new_table([
+    string_col("${REQUIREMENTS_TABLE_NAME_COLUMN_NAME}", list(installed)),
+    string_col("${REQUIREMENTS_TABLE_VERSION_COLUMN_NAME}", [version(pkg) for pkg in installed])
+])` as const;

src/controllers/ExtensionController.ts

@@ -5,6 +5,7 @@ import {
   CLEAR_SECRET_STORAGE_CMD,
   CREATE_NEW_TEXT_DOC_CMD,
   DOWNLOAD_LOGS_CMD,
+  GENERATE_REQUIREMENTS_TXT_CMD,
   OPEN_IN_BROWSER_CMD,
   REFRESH_SERVER_CONNECTION_TREE_CMD,
   REFRESH_SERVER_TREE_CMD,
@@ -66,6 +67,7 @@ import type {
   CoreAuthenticatedClient,
   ICoreClientFactory,
   CoreUnauthenticatedClient,
+  ConnectionState,
 } from '../types';
 import { ServerConnectionTreeDragAndDropController } from './ServerConnectionTreeDragAndDropController';
 import { ConnectionController } from './ConnectionController';
@@ -454,6 +456,12 @@ export class ExtensionController implements Disposable {
     /** Download logs and open in editor */
     this.registerCommand(DOWNLOAD_LOGS_CMD, this.onDownloadLogs);
 
+    /** Generate requirements.txt */
+    this.registerCommand(
+      GENERATE_REQUIREMENTS_TXT_CMD,
+      this.onGenerateRequirementsTxt
+    );
+
     /** Open server in browser */
     this.registerCommand(OPEN_IN_BROWSER_CMD, this.onOpenInBrowser);
 
@@ -636,6 +644,19 @@ export class ExtensionController implements Disposable {
     }
   };
 
+  /**
+   * Handle generating requirements.txt command
+   */
+  onGenerateRequirementsTxt = async (
+    connectionState: ConnectionState
+  ): Promise<void> => {
+    if (!isInstanceOf(connectionState, DhcService)) {
+      throw new Error('Connection is not a DHC service');
+    }
+
+    await connectionState.generateRequirementsTxt();
+  };
+
   onOpenInBrowser = async (serverState: ServerState): Promise<void> => {
     const psk = await this._secretService?.getPsk(serverState.url);
     const serverUrl = new URL(serverState.url);

src/dh/dhc.ts

@@ -3,8 +3,16 @@ import { NoConsoleTypesError } from './errorUtils';
 import type {
   CoreAuthenticatedClient,
   CoreUnauthenticatedClient,
+  DependencyName,
+  DependencyVersion,
 } from '../types';
 import { hasStatusCode, loadModules } from '@deephaven/jsapi-nodejs';
+import {
+  REQUIREMENTS_QUERY_TXT,
+  REQUIREMENTS_TABLE_NAME,
+  REQUIREMENTS_TABLE_NAME_COLUMN_NAME,
+  REQUIREMENTS_TABLE_VERSION_COLUMN_NAME,
+} from '../common';
 
 export const AUTH_HANDLER_TYPE_ANONYMOUS =
   'io.deephaven.auth.AnonymousAuthenticationHandler';
@@ -89,6 +97,38 @@ export function getEmbedWidgetUrl({
   return url;
 }
 
+/**
+ * Get a name / version map of python dependencies from a DH session.
+ * @param session The DH session to use.
+ * @returns A promise that resolves to a map of python dependencies.
+ */
+export async function getPythonDependencies(
+  session: DhType.IdeSession
+): Promise<Map<DependencyName, DependencyVersion>> {
+  await session.runCode(REQUIREMENTS_QUERY_TXT);
+
+  const dependencies = new Map<DependencyName, DependencyVersion>();
+
+  const table = await session.getTable(REQUIREMENTS_TABLE_NAME);
+  table.setViewport(0, table.size - 1);
+  const data = await table.getViewportData();
+
+  const nameColumn = table.findColumn(REQUIREMENTS_TABLE_NAME_COLUMN_NAME);
+  const versionColumn = table.findColumn(
+    REQUIREMENTS_TABLE_VERSION_COLUMN_NAME
+  );
+
+  for (const row of data.rows) {
+    const name: DependencyName = row.get(nameColumn);
+    const version: DependencyVersion = row.get(versionColumn);
+    dependencies.set(name, version);
+  }
+
+  await session.runCode(`del ${REQUIREMENTS_TABLE_NAME}`);
+
+  return dependencies;
+}
+
 /**
  * Initialize a connection and session to a DHC server.
  * @param client The authenticated client to use for the connection.

src/providers/ServerConnectionTreeProvider.ts

@@ -20,6 +20,7 @@ export class ServerConnectionTreeProvider extends TreeDataProviderBase<ServerCon
     if (connectionOrUri instanceof vscode.Uri) {
       return {
         description: connectionOrUri.path,
+        contextValue: CONNECTION_TREE_ITEM_CONTEXT.isUri,
         command: {
           command: 'vscode.open',
           title: 'Open Uri',
@@ -58,7 +59,9 @@ export class ServerConnectionTreeProvider extends TreeDataProviderBase<ServerCon
     return {
       label,
       description: descriptionTokens.join(' - '),
-      contextValue: CONNECTION_TREE_ITEM_CONTEXT.isConnection,
+      contextValue: connectionOrUri.isConnected
+        ? CONNECTION_TREE_ITEM_CONTEXT.isConnectionConnected
+        : CONNECTION_TREE_ITEM_CONTEXT.isConnectionConnecting,
       collapsibleState: hasUris
         ? vscode.TreeItemCollapsibleState.Expanded
         : undefined,

src/services/DhcService.ts

@@ -1,8 +1,18 @@
 import * as vscode from 'vscode';
 import { isAggregateError } from '@deephaven/jsapi-nodejs';
 import type { dh as DhcType } from '@deephaven/jsapi-types';
-import { formatTimestamp, getCombinedSelectedLinesText, Logger } from '../util';
-import { initDhcSession, type ConnectionAndSession } from '../dh/dhc';
+import {
+  assertDefined,
+  formatTimestamp,
+  getCombinedSelectedLinesText,
+  Logger,
+  saveRequirementsTxt,
+} from '../util';
+import {
+  getPythonDependencies,
+  initDhcSession,
+  type ConnectionAndSession,
+} from '../dh/dhc';
 import type {
   ConsoleType,
   CoreAuthenticatedClient,
@@ -298,6 +308,22 @@ export class DhcService implements IDhcService {
     return consoleTypes.has(consoleType);
   };
 
+  /**
+   * Generate a requirements.txt file based on the packages used in the current
+   * session.
+   */
+  async generateRequirementsTxt(): Promise<void> {
+    if (this.session == null) {
+      await this.initSession();
+    }
+
+    assertDefined(this.session, 'session');
+
+    const dependencies = await getPythonDependencies(this.session);
+
+    await saveRequirementsTxt(dependencies);
+  }
+
   async runEditorCode(
     editor: vscode.TextEditor,
     selectionOnly = false

src/types/commonTypes.d.ts

@@ -45,6 +45,9 @@ export type CoreUnauthenticatedClient = Brand<
   DhcType.CoreClient
 >;
 
+export type DependencyName = Brand<'DependencyName', string>;
+export type DependencyVersion = Brand<'DependencyVersion', string>;
+
 export type EnterpriseConnectionConfigStored =
   | Brand<'EnterpriseConnectionConfigStored'>
   | { url: string; label?: string; experimentalWorkerConfig?: WorkerConfig };

src/util/uiUtils.ts

@@ -1,4 +1,5 @@
 import * as vscode from 'vscode';
+import * as fs from 'node:fs';
 import {
   SELECT_CONNECTION_COMMAND,
   STATUS_BAR_CONNECTING_TEXT,
@@ -15,6 +16,8 @@ import type {
   ConnectionState,
   UserLoginPreferences,
   Psk,
+  DependencyName,
+  DependencyVersion,
 } from '../types';
 import { sortByStringProp } from './dataUtils';
 import { assertDefined } from './assertUtil';
@@ -358,6 +361,28 @@ export async function getEditorForUri(
   return vscode.window.showTextDocument(uri, { preview: false, viewColumn });
 }
 
+/**
+ * Get the workspace folder for the active editor or fallback to the first listed
+ * workspace folder.
+ * @returns The workspace folder or undefined if there are no workspace folders.
+ */
+export function getWorkspaceFolder(): vscode.WorkspaceFolder | undefined {
+  const wkspFolders = vscode.workspace.workspaceFolders ?? [];
+
+  if (wkspFolders.length === 0) {
+    return;
+  }
+
+  const activeUri = vscode.window.activeTextEditor?.document.uri;
+
+  const activeWkspFolder =
+    activeUri == null
+      ? null
+      : wkspFolders.find(path => activeUri.fsPath.startsWith(path.uri.fsPath));
+
+  return activeWkspFolder ?? wkspFolders[0];
+}
+
 /**
  * Update given status bar item based on connection status
  * and optional `ConnectionOption`.
@@ -445,3 +470,39 @@ export function promptForOperateAs(
     value: defaultValue,
   }) as Promise<OperateAsUsername | undefined>;
 }
+
+/**
+ * Save a map of dependency name / versions to a `requirements.txt` file.
+ * @param dependencies The map of dependency names / versions to save.
+ * @returns Promise that resolves when the file is saved.
+ */
+export async function saveRequirementsTxt(
+  dependencies: Map<DependencyName, DependencyVersion>
+): Promise<void> {
+  const wkspFolder = getWorkspaceFolder();
+
+  const defaultUri =
+    wkspFolder == null
+      ? vscode.Uri.file('requirements.txt')
+      : vscode.Uri.joinPath(wkspFolder.uri, 'requirements.txt');
+
+  const uri = await vscode.window.showSaveDialog({
+    defaultUri,
+    // eslint-disable-next-line @typescript-eslint/naming-convention
+    filters: { Requirements: ['txt'] },
+  });
+
+  if (uri == null) {
+    return;
+  }
+
+  const sorted = [
+    ...dependencies
+      .entries()
+      .map(([packageName, version]) => `${packageName}==${version}`),
+  ].sort((a, b) => a.localeCompare(b));
+
+  fs.writeFileSync(uri.fsPath, sorted.join('\n'));
+
+  vscode.window.showTextDocument(uri);
+}