Set notebook cell language back to Python if %%qsharp magic isn't there

vscode/src/notebook.ts

@@ -47,24 +47,49 @@ export function registerQSharpNotebookHandlers() {
         const addedCells = event.contentChanges
           .map((change) => change.addedCells)
           .flat();
+
         updateQSharpCellLanguages(changedCells.concat(addedCells));
       }
     }),
   );
 
+  let defaultLanguageId: string | undefined;
+
   function updateQSharpCellLanguages(cells: vscode.NotebookCell[]) {
     for (const cell of cells) {
-      // If this is a code cell that starts with %%qsharp, and language isn't already set to Q#, set it.
+      // If this is a code cell that starts with %%qsharp, and language wasn't already set to Q#, set it.
       if (cell.kind === vscode.NotebookCellKind.Code) {
         const document = cell.document;
-        if (
-          document.languageId !== qsharpLanguageId &&
-          findQSharpCellMagic(document)
-        ) {
-          vscode.languages.setTextDocumentLanguage(
-            cell.document,
-            qsharpLanguageId,
-          );
+        const currentLanguageId = document.languageId;
+        if (findQSharpCellMagic(document)) {
+          if (currentLanguageId !== qsharpLanguageId) {
+            // Remember the "default" language of the notebook (this will normally be Python)
+            defaultLanguageId = currentLanguageId;
+            vscode.languages.setTextDocumentLanguage(
+              cell.document,
+              qsharpLanguageId,
+            );
+            log.trace(
+              `setting cell ${cell.index} language to ${qsharpLanguageId}`,
+            );
+          }
+        } else {
+          // This is not a %%qsharp cell. If the language was set to Q#,
+          // change it back to the default language.
+          //
+          // If the cell language was not set to Q#, it's out of our purview and we don't
+          // want to automatically change the language settings. For example, this could
+          // be a %%bash cell magic and the user may have intentionally set the language
+          // to "shell".
+          if (currentLanguageId === qsharpLanguageId && defaultLanguageId) {
+            vscode.languages.setTextDocumentLanguage(
+              cell.document,
+              defaultLanguageId,
+            );
+            log.trace(
+              `setting cell ${cell.index} language to ${defaultLanguageId}`,
+            );
+          }
         }
       }
     }

vscode/test/suites/extensionUtils.ts

@@ -31,3 +31,66 @@ export async function activateExtension() {
     `qsharp-tests: activate() completed in ${performance.now() - start}ms`,
   );
 }
+
+/**
+ * Waits until a condition is met, with a timeout.
+ *
+ * @param condition A function that returns true when the condition is met
+ * @param wakeUpOn The event that we want to wake up to check the condition
+ * @param timeoutMs If the condition is not met by this timeout, this function will throw
+ * @param timeoutErrorMsg The custom error message to throw if the condition is not met
+ */
+export async function waitForCondition(
+  condition: (...args: any[]) => boolean,
+  wakeUpOn: vscode.Event<any>,
+  timeoutMs: number,
+  timeoutErrorMsg: string,
+) {
+  let disposable: vscode.Disposable | undefined;
+  await new Promise<void>((resolve, reject) => {
+    let done = false;
+    setTimeout(() => {
+      if (!done) {
+        reject(new Error(timeoutErrorMsg));
+      }
+    }, timeoutMs);
+
+    disposable = wakeUpOn(() => {
+      if (!done && condition()) {
+        done = true;
+        resolve();
+      }
+    });
+
+    // Resolve immediately if condition is already met
+    if (condition()) {
+      done = true;
+      resolve();
+    }
+  });
+  disposable?.dispose();
+}
+
+/**
+ * Convenience method to add a time delay.
+ *
+ * @deprecated While this function is convenient for local development,
+ * please do NOT use it in actual tests. Instead, use `waitForCondition`
+ * with an appropriate condition and wake up event.
+ *
+ * Adding an unconditional delay to tests is guaranteed to slow tests down
+ * unnecessarily. It also reduces reliability when we're not clear about
+ * exactly what we're waiting on.
+ */
+export async function delay(timeoutMs: number) {
+  try {
+    await waitForCondition(
+      () => false,
+      () => ({ dispose() {} }),
+      timeoutMs,
+      "hit the expected timeout",
+    );
+  } catch (e) {
+    // expected
+  }
+}

vscode/test/suites/language-service/notebook.test.ts

@@ -3,7 +3,7 @@
 
 import * as vscode from "vscode";
 import { assert } from "chai";
-import { activateExtension } from "../extensionUtils";
+import { activateExtension, waitForCondition } from "../extensionUtils";
 
 suite("Q# Notebook Tests", function suite() {
   const workspaceFolder =
@@ -23,40 +23,56 @@ suite("Q# Notebook Tests", function suite() {
 
     // Test if the cell with the %%qsharp magic has been detected
     // and its language switched to qsharp. We can only expect this to happen
-    // after the notebook has happened, and the document handlers
+    // after the notebook has been opened, and the document handlers
     // have been invoked, but of course there is no callback for that.
     // So we just verify the notebook has been updated with the qsharp
     // language id within 50ms (If we start exceeding this timeout for some
     // reason, it's enough of a user-perceptible delay that we're probably
     // better off disabling this behavior, rather than suddenly change the
     // cell language from under the user after a delay).
-    await new Promise<void>((resolve, reject) => {
-      let done = false;
-      setTimeout(() => {
-        if (!done) {
-          reject(new Error("timed out waiting for a Q# code cell"));
-        }
-      }, 50);
-
-      vscode.workspace.onDidChangeNotebookDocument((event) => {
-        if (!done && hasQSharpCell(event.notebook)) {
-          done = true;
-          resolve();
-        }
-      });
-
-      // in case the notebook updates have already occurred by the time we get here
-      if (hasQSharpCell(notebook)) {
-        done = true;
-        resolve();
-      }
-
-      function hasQSharpCell(notebook) {
-        return notebook
+    await waitForCondition(
+      () =>
+        !!notebook
           .getCells()
-          .find((cell) => cell.document.languageId === "qsharp");
-      }
-    });
+          .find((cell) => cell.document.languageId === "qsharp"),
+      vscode.workspace.onDidChangeNotebookDocument,
+      50,
+      "timed out waiting for a Q# code cell",
+    );
+  });
+
+  test("Cell language is set back to Python", async () => {
+    const notebook = await vscode.workspace.openNotebookDocument(
+      vscode.Uri.joinPath(workspaceFolderUri, "test.ipynb"),
+    );
+
+    await vscode.window.showNotebookDocument(notebook);
+
+    assert.equal(
+      vscode.window.activeNotebookEditor?.notebook.uri.toString(),
+      notebook.uri.toString(),
+    );
+
+    const oldLength = notebook.getCells().length;
+
+    // Add a new cell at the bottom of the notebook
+    await vscode.commands.executeCommand("notebook.focusBottom");
+    await vscode.commands.executeCommand("notebook.cell.insertCodeCellBelow");
+
+    // There should be an additional cell in the notebook and it should be Python
+    await waitForCondition(
+      () => {
+        const cellsAfter = notebook.getCells();
+        return (
+          cellsAfter.length === oldLength + 1 &&
+          notebook.getCells()[cellsAfter.length - 1].document.languageId ===
+            "python"
+        );
+      },
+      vscode.workspace.onDidChangeNotebookDocument,
+      50,
+      "timed out waiting for a Python code cell",
+    );
   });
 
   test("Diagnostics", async () => {