Support multi-root Pydance workspaces

pydance/package.json

@@ -2,7 +2,7 @@
     "name": "pydance",
     "displayName": "Pydance",
     "description": "Python language server that provides high-performance workspace symbol search for large Python codebases",
-    "version": "0.3.0",
+    "version": "0.3.1",
     "publisher": "ToughType",
     "icon": "images/pydance-logo.png",
     "repository": {
@@ -20,6 +20,12 @@
     ],
     "main": "./out/extension.js",
     "contributes": {
+        "commands": [
+            {
+                "command": "pydance.restartServer",
+                "title": "Pydance: Restart Server"
+            }
+        ],
         "configuration": {
             "title": "Pydance",
             "properties": {
@@ -62,6 +68,7 @@
         "pretest": "npm run compile:tests",
         "test": "node ./out/test/runTest.js",
         "test:integration": "npm run compile:tests && node ./out/test/runIntegrationTest.js",
+        "test:integration:multi-root": "npm run compile:tests && node ./out/test/runMultiRootIntegrationTest.js",
         "lint": "eslint src --ext ts",
         "lint:fix": "eslint src --ext ts --fix",
         "format": "prettier --write \"src/**/*.ts\"",

pydance/src/extension.ts

@@ -7,7 +7,19 @@ import {
   Trace,
 } from "vscode-languageclient/node";
 
-let client: LanguageClient;
+let client: LanguageClient | undefined;
+let serverOptions: ServerOptions;
+let clientOptions: LanguageClientOptions;
+let trace: Trace = Trace.Off;
+
+async function restartServer() {
+  if (client) {
+    await client.stop();
+  }
+  client = new LanguageClient("pydance", "Pydance", serverOptions, clientOptions);
+  client.setTrace(trace);
+  await client.start();
+}
 
 export function activate(context: vscode.ExtensionContext) {
   const outputChannel = vscode.window.createOutputChannel("Pydance");
@@ -23,13 +35,13 @@ export function activate(context: vscode.ExtensionContext) {
   const excludePatterns = config.get<string[]>("excludePatterns", []);
   outputChannel.appendLine(`Using parser: ${parser}`);
   // If the extension is launched in debug mode then the debug server options are used
-  const serverOptions: ServerOptions = {
+  serverOptions = {
     run: { command: serverPath, args: ["--parser", parser] },
     debug: { command: serverPath, args: ["--parser", parser] },
   };
 
   // Options to control the language client
-  const clientOptions: LanguageClientOptions = {
+  clientOptions = {
     // Register the server for Python documents
     documentSelector: [{ scheme: "file", language: "python" }],
     outputChannel: outputChannel,
@@ -40,21 +52,20 @@ export function activate(context: vscode.ExtensionContext) {
   };
 
   // Create the language client and start the client.
-  client = new LanguageClient(
-    "pydance",
-    "Pydance",
-    // serverOptions,
-    serverOptions,
-    clientOptions
-  );
+  client = new LanguageClient("pydance", "Pydance", serverOptions, clientOptions);
 
   // Set trace level based on configuration
   const traceMap: { [key: string]: Trace } = {
     off: Trace.Off,
     messages: Trace.Messages,
     verbose: Trace.Verbose,
   };
-  client.setTrace(traceMap[traceLevel] || Trace.Off);
+  trace = traceMap[traceLevel] || Trace.Off;
+  client.setTrace(trace);
+
+  context.subscriptions.push(
+    vscode.commands.registerCommand("pydance.restartServer", restartServer)
+  );
 
   outputChannel.appendLine("Starting language client...");
   // Start the client. This will also launch the server

pydance/src/test/runIntegrationTest.ts

@@ -1,3 +1,4 @@
+import * as fs from "fs";
 import * as path from "path";
 import { runTests } from "@vscode/test-electron";
 
@@ -12,13 +13,16 @@ async function main() {
     // The path to the test workspace - this ensures we have a proper workspace
     const testWorkspace = path.resolve(__dirname, "../../src/testFixture");
 
+    const userDataDir = path.join("/tmp", `pydance-vscode-test-${process.pid}`);
+    fs.rmSync(userDataDir, { recursive: true, force: true });
+
     console.log("Running integration tests with workspace:", testWorkspace);
 
     // Download VS Code, unzip it and run the integration test
     await runTests({
       extensionDevelopmentPath,
       extensionTestsPath,
-      launchArgs: [testWorkspace],
+      launchArgs: [testWorkspace, "--user-data-dir", userDataDir],
       // Set environment variable to indicate integration test mode
       extensionTestsEnv: {
         INTEGRATION_TEST: "true",

pydance/src/test/runMultiRootIntegrationTest.ts

@@ -0,0 +1,32 @@
+import * as fs from "fs";
+import * as path from "path";
+import { runTests } from "@vscode/test-electron";
+
+async function main() {
+  try {
+    const extensionDevelopmentPath = path.resolve(__dirname, "../../");
+    const extensionTestsPath = path.resolve(__dirname, "./suite/index");
+    const testWorkspace = path.resolve(
+      __dirname,
+      "../../src/testWorkspaces/multi-root.code-workspace"
+    );
+    const userDataDir = path.join("/tmp", `pydance-vscode-test-${process.pid}`);
+    fs.rmSync(userDataDir, { recursive: true, force: true });
+
+    console.log("Running multi-root integration tests with workspace:", testWorkspace);
+
+    await runTests({
+      extensionDevelopmentPath,
+      extensionTestsPath,
+      launchArgs: [testWorkspace, "--user-data-dir", userDataDir],
+      extensionTestsEnv: {
+        INTEGRATION_TEST: "true",
+      },
+    });
+  } catch (err) {
+    console.error("Failed to run multi-root integration tests");
+    process.exit(1);
+  }
+}
+
+main();

pydance/src/test/suite/extension.test.ts

@@ -82,34 +82,11 @@ suite("Extension Test Suite", () => {
   });
 
   suite("Integration Tests", () => {
-    test("Should provide workspace symbols with real language server", async function () {
-      // Skip if not running in integration test mode
-      if (process.env.INTEGRATION_TEST !== "true") {
-        console.log("Not in integration test mode, skipping");
-        this.skip();
-      }
-
-      // First check if we have a workspace folder
-      const workspaceFolders = vscode.workspace.workspaceFolders;
-      if (!workspaceFolders || workspaceFolders.length === 0) {
-        console.log("No workspace folder found, skipping integration test");
-        this.skip();
-      }
-
-      // Check if pylight binary exists
-      // eslint-disable-next-line @typescript-eslint/no-var-requires
-      const fs = require("fs");
-      // eslint-disable-next-line @typescript-eslint/no-var-requires
-      const path = require("path");
-      const extensionPath =
-        vscode.extensions.getExtension("ToughType.pydance")!.extensionPath;
-      const pylightPath = path.join(extensionPath, "pylight");
-
-      if (!fs.existsSync(pylightPath)) {
-        console.log("pylight binary not found, skipping integration test");
-        this.skip();
-      }
+    setup(function () {
+      skipUnlessIntegrationReady(this);
+    });
 
+    test("Should provide workspace symbols with real language server", async function () {
       await testWorkspaceSymbols(docUri, [
         new vscode.SymbolInformation(
           "TestClass",
@@ -173,9 +150,100 @@ suite("Extension Test Suite", () => {
         ),
       ]);
     });
+
+    test("Should index all folders in a multi-root workspace", async function () {
+      const repoTwoFolder = secondWorkspaceFolder();
+      if (!repoTwoFolder) {
+        console.log("No second workspace folder found, skipping multi-root test");
+        this.skip();
+      }
+
+      await activate(docUri);
+      const symbols = await waitForWorkspaceSymbols("repo_two", 2);
+      const names = symbols.map((symbol) => symbol.name);
+
+      assert.ok(
+        names.includes("repo_two_function"),
+        "repo_two_function should be indexed"
+      );
+      assert.ok(
+        names.includes("repo_two_helper"),
+        "repo_two_helper should be indexed"
+      );
+      assert.ok(
+        symbols.every((symbol) =>
+          symbol.location.uri.fsPath.startsWith(repoTwoFolder!.uri.fsPath)
+        ),
+        "repo_two results should come from the second workspace folder"
+      );
+    });
+
+    test("Should expose a restart command that reindexes", async function () {
+      await activate(docUri);
+      await vscode.commands.executeCommand("pydance.restartServer");
+
+      const symbols = await waitForWorkspaceSymbols("test", 4);
+      assert.ok(
+        symbols.some((symbol) => symbol.name === "test_function"),
+        "Expected test_function after restart"
+      );
+    });
   });
 });
 
+function secondWorkspaceFolder() {
+  return vscode.workspace.workspaceFolders?.find(
+    (folder) => folder.name === "testFixtureSecond"
+  );
+}
+
+function skipUnlessIntegrationReady(context: Mocha.Context) {
+  if (process.env.INTEGRATION_TEST !== "true") {
+    console.log("Not in integration test mode, skipping");
+    context.skip();
+  }
+
+  const workspaceFolders = vscode.workspace.workspaceFolders;
+  if (!workspaceFolders || workspaceFolders.length === 0) {
+    console.log("No workspace folder found, skipping integration test");
+    context.skip();
+  }
+
+  // Check if pylight binary exists
+  // eslint-disable-next-line @typescript-eslint/no-var-requires
+  const fs = require("fs");
+  // eslint-disable-next-line @typescript-eslint/no-var-requires
+  const path = require("path");
+  const extensionPath =
+    vscode.extensions.getExtension("ToughType.pydance")!.extensionPath;
+  const pylightPath = path.join(extensionPath, "pylight");
+
+  if (!fs.existsSync(pylightPath)) {
+    console.log("pylight binary not found, skipping integration test");
+    context.skip();
+  }
+}
+
+async function waitForWorkspaceSymbols(query: string, minCount: number) {
+  const deadline = Date.now() + 10_000;
+  let symbols: vscode.SymbolInformation[] = [];
+
+  while (Date.now() < deadline) {
+    symbols = await vscode.commands.executeCommand<vscode.SymbolInformation[]>(
+      "vscode.executeWorkspaceSymbolProvider",
+      query
+    );
+    if (symbols.length >= minCount) {
+      return symbols;
+    }
+    await new Promise((resolve) => setTimeout(resolve, 250));
+  }
+
+  assert.fail(
+    `Expected at least ${minCount} symbols for ${query}, found ${symbols.length}`
+  );
+}
+
 async function testWorkspaceSymbols(
   docUri: vscode.Uri,
   expectedSymbols: vscode.SymbolInformation[]

pydance/src/testFixtureSecond/second.py

@@ -0,0 +1,11 @@
+class RepoTwoClass:
+    def repo_two_method(self):
+        pass
+
+
+def repo_two_function():
+    return "repo two"
+
+
+def repo_two_helper():
+    return "helper"

pydance/src/testWorkspaces/multi-root.code-workspace

@@ -0,0 +1,6 @@
+{
+  "folders": [
+    { "path": "../testFixture" },
+    { "path": "../testFixtureSecond" }
+  ]
+}

pylight/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pylight"
-version = "0.3.0"
+version = "0.3.1"
 edition = "2021"