feat: improve handling of prepend

language-configuration.json

@@ -25,8 +25,8 @@
   ],
   "folding": {
     "markers": {
-      "start": "// PREPEND -- DO NOT EDIT",
-      "end": "// END PREPEND"
+      "start": "// PREPEND",
+      "end": "// END PREPEND -- DO NOT EDIT PREPEND"
     }
   },
   "wordPattern": "[_$a-zA-Z][_$a-zA-Z0-9]*",

src/commands/showPanel.tsx

@@ -7,10 +7,11 @@ import { LANGUAGES } from "../utils/languages";
 import { setWebviewContent } from "../utils/webview";
 import config from "../utils/config";
 import { FRONTEND_ELEMENT_ID } from "../constants";
+import _ from "lodash";
 import { MessageHandler } from "../utils/messageHandler";
+import { SOURCE_ACADEMY_ICON_URI } from "../extension";
 
 let messageHandler = MessageHandler.getInstance();
-import { SOURCE_ACADEMY_ICON_URI } from "../extension";
 
 export async function showPanel(
   context: vscode.ExtensionContext,

src/utils/editor.ts

@@ -1,10 +1,12 @@
+// TODO: Move this file to src/features/editor/editor.ts
 import * as vscode from "vscode";
 
 import config from "../utils/config";
 import Messages, { VscWorkspaceLocation } from "./messages";
 import path from "path";
 import { sendToFrontendWrapped } from "../commands/showPanel";
 import { canonicaliseLocation } from "./misc";
+import { codeAddPrepend, codeRemovePrepend } from "./editorUtils";
 
 export class Editor {
   editor?: vscode.TextEditor;
@@ -62,58 +64,53 @@ export class Editor {
     const uri = vscode.Uri.file(filePath);
     self.uri = uri.toString();
 
-    const contents =
-      prepend !== ""
-        ? [
-            "// PREPEND -- DO NOT EDIT",
-            prepend,
-            "// END PREPEND",
-            initialCode,
-          ].join("\n")
-        : initialCode;
-
-    await vscode.workspace.fs.readFile(vscode.Uri.file(filePath)).then(
-      (value) => {
-        if (value.toString() !== contents) {
-          self.log(
-            "EXTENSION: Conflict detected between local and remote, prompting user to choose one",
+    const contents = codeAddPrepend(initialCode, prepend);
+
+    await vscode.workspace.fs
+      .readFile(vscode.Uri.file(filePath))
+      .then((value) => value.toString())
+      .then(
+        (localCode) => {
+          if (localCode !== contents) {
+            self.log(
+              "EXTENSION: Conflict detected between local and remote, prompting user to choose one",
+            );
+            vscode.window
+              .showInformationMessage(
+                [
+                  "The local file differs from the version on the Source Academy servers.",
+                  "Discard the local file and use the one on the server?",
+                ].join(" "),
+                { modal: true },
+                "Yes",
+              )
+              .then(async (answer) => {
+                // By default the code displayed is the local one
+                if (answer === "Yes") {
+                  self.log("EXTENSION: Saving program from server to file");
+                  await vscode.workspace.fs.writeFile(
+                    uri,
+                    new TextEncoder().encode(contents),
+                  );
+                } else if (answer === undefined) {
+                  // Modal cancelled
+                  const message = Messages.Text(
+                    self.workspaceLocation,
+                    codeRemovePrepend(localCode),
+                  );
+                  sendToFrontendWrapped(message);
+                }
+              });
+          }
+        },
+        async () => {
+          self.log(`Opening file failed, creating at ${filePath}`);
+          await vscode.workspace.fs.writeFile(
+            uri,
+            new TextEncoder().encode(contents),
           );
-          vscode.window
-            .showInformationMessage(
-              [
-                "The local file differs from the version on the Source Academy servers.",
-                "Discard the local file and use the one on the server?",
-              ].join(" "),
-              { modal: true },
-              "Yes",
-            )
-            .then(async (answer) => {
-              // By default the code displayed is the local one
-              if (answer === "Yes") {
-                self.log("EXTENSION: Saving program from server to file");
-                await vscode.workspace.fs.writeFile(
-                  uri,
-                  new TextEncoder().encode(contents),
-                );
-              } else if (answer === undefined) {
-                // Modal cancelled
-                const message = Messages.Text(
-                  self.workspaceLocation,
-                  value.toString(),
-                );
-                sendToFrontendWrapped(message);
-              }
-            });
-        }
-      },
-      async () => {
-        self.log(`Opening file failed, creating at ${filePath}`);
-        await vscode.workspace.fs.writeFile(
-          uri,
-          new TextEncoder().encode(contents),
-        );
-      },
-    );
+        },
+      );
 
     const editor = await vscode.window.showTextDocument(uri, {
       preview: false,

src/utils/editorUtils.ts

@@ -0,0 +1,44 @@
+// TODO: Move this file to src/features/editor/utils.ts
+
+// ================================================================================
+// Prepend
+// ================================================================================
+
+// If this is edited, also change folding.markers.{start,end} in language-configuration.json
+const PREPEND_MARKER_START = "// PREPEND";
+const PREPEND_MARKER_END = "// END PREPEND -- DO NOT EDIT PREPEND";
+
+export function codeAddPrepend(code: string, prepend: string) {
+  if (prepend === "") {
+    return code;
+  }
+  return [
+    PREPEND_MARKER_START,
+    prepend,
+    PREPEND_MARKER_END,
+    "", // Visual separation between prepend and code
+    code,
+  ].join("\n");
+}
+
+export function getNumPrependLines(prepend: string) {
+  if (prepend === "") {
+    return 0;
+  }
+  return prepend.split("\n").length + 3; // account for start/end markers
+}
+
+export function codeRemovePrepend(code: string) {
+  const lines = code.split("\n");
+  const start = lines.indexOf(PREPEND_MARKER_START);
+  const end = lines.indexOf(PREPEND_MARKER_END);
+
+  if (start === -1 || end === -1 || end < start) {
+    return code;
+  }
+
+  // If line spacing between prepend and code accidentally removed, do not delete code!
+  const skip = end + (lines[end + 1] === "" ? 1 : 0);
+
+  return lines.slice(skip + 1).join("\n");
+}

src/utils/messageHandler.tsx

@@ -14,6 +14,7 @@ import { client } from "../extension";
 import _ from "lodash";
 import { treeDataProvider } from "../treeview";
 import McqPanel from "../webview/components/McqPanel";
+import { codeRemovePrepend, getNumPrependLines } from "./editorUtils";
 
 /*
  * MessageHandler is a singleton class that handles messages from the frontend
@@ -122,18 +123,12 @@ export class MessageHandler {
           this.activeEditor.uri;
           const info = context.globalState.get("info") ?? {};
           if (this.activeEditor.uri) {
-            // TODO: Write our own wrapper to set nested keys easily, removing lodash
-            // @ts-ignore
             _.set(
               info,
               `["${this.activeEditor.uri}"].chapter`,
               message.chapter ?? 1,
             );
-            // TODO: message.prepend can be undefined in runtime, investigate
-            const nPrependLines =
-              message.prepend && message.prepend !== ""
-                ? message.prepend.split("\n").length + 2 // account for start/end markers
-                : 0;
+            const nPrependLines = getNumPrependLines(message.prepend);
             _.set(info, `["${this.activeEditor.uri}"].prepend`, nPrependLines);
             context.globalState.update("info", info);
             client.sendRequest("source/publishInfo", info);
@@ -155,7 +150,10 @@ export class MessageHandler {
                 `EXTENSION: Editor ${editor.assessmentName}_${editor.questionId}_${editor.assessmentType} is no longer active, skipping onChange`,
               );
             }
-            const message = Messages.Text(workspaceLocation, code);
+            const message = Messages.Text(
+              workspaceLocation,
+              codeRemovePrepend(code),
+            );
             console.log(`Sending message: ${JSON.stringify(message)}`);
             sendToFrontend(this.panel, message);
           });