Extracted method for getting suggestion marks.

.yarn/versions/24a8541c.yml

@@ -0,0 +1,2 @@
+releases:
+  "@handlewithcare/prosemirror-suggest-changes": patch

src/addNodeMarkStep.ts

@@ -3,6 +3,7 @@ import { type EditorState, type Transaction } from "prosemirror-state";
 import { type AddNodeMarkStep, type Step } from "prosemirror-transform";
 
 import { rebasePos } from "./rebasePos.js";
+import { getSuggestionMarks } from "./utils.js";
 
 /**
  * Transform an add node mark step into its equivalent tracked steps.
@@ -18,12 +19,7 @@ export function trackAddNodeMarkStep(
   prevSteps: Step[],
   suggestionId: number,
 ) {
-  const { modification } = state.schema.marks;
-  if (!modification) {
-    throw new Error(
-      `Failed to apply modifications to node: schema does not contain modification mark. Did you forget to add it?`,
-    );
-  }
+  const { modification } = getSuggestionMarks(state.schema);
 
   const rebasedPos = rebasePos(step.pos, prevSteps, trackedTransaction.steps);
   const $pos = trackedTransaction.doc.resolve(rebasedPos);

src/attrStep.ts

@@ -3,6 +3,7 @@ import { type EditorState, type Transaction } from "prosemirror-state";
 import { type AttrStep, type Step } from "prosemirror-transform";
 
 import { rebasePos } from "./rebasePos.js";
+import { getSuggestionMarks } from "./utils.js";
 
 /**
  * Transform an attr mark step into its equivalent tracked steps.
@@ -18,12 +19,7 @@ export function trackAttrStep(
   prevSteps: Step[],
   suggestionId: number,
 ) {
-  const { modification } = state.schema.marks;
-  if (!modification) {
-    throw new Error(
-      `Failed to apply modifications to node: schema does not contain modification mark. Did you forget to add it?`,
-    );
-  }
+  const { modification } = getSuggestionMarks(state.schema);
 
   const rebasedPos = rebasePos(step.pos, prevSteps, trackedTransaction.steps);
   const $pos = trackedTransaction.doc.resolve(rebasedPos);

src/commands.ts

@@ -15,6 +15,7 @@ import { type EditorView } from "prosemirror-view";
 
 import { findSuggestionMarkEnd } from "./findSuggestionMarkEnd.js";
 import { suggestChangesKey } from "./plugin.js";
+import { getSuggestionMarks } from "./utils.js";
 
 /**
  * Given a node and a transform, add a set of steps to the
@@ -103,9 +104,8 @@ function applySuggestionsToTransform(
 }
 
 function revertModifications(node: Node, pos: number, tr: Transform) {
-  const existingMods = node.marks.filter(
-    (mark) => mark.type === node.type.schema.marks["modification"],
-  );
+  const { modification } = getSuggestionMarks(node.type.schema);
+  const existingMods = node.marks.filter((mark) => mark.type === modification);
   for (const mod of existingMods) {
     if (
       mod.attrs["type"] === "attr" &&
@@ -146,12 +146,7 @@ function applyModificationsToTransform(
   dir: number,
   suggestionId?: number,
 ) {
-  const { modification } = node.type.schema.marks;
-  if (!modification) {
-    throw new Error(
-      `Failed to apply modifications to node: schema does not contain modification mark. Did you forget to add it?`,
-    );
-  }
+  const { modification } = getSuggestionMarks(node.type.schema);
 
   const modificationIsInSet =
     suggestionId === undefined
@@ -194,17 +189,7 @@ function applyModificationsToTransform(
 }
 
 function applySuggestionsToNode(node: Node) {
-  const { deletion, insertion } = node.type.schema.marks;
-  if (!deletion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain deletion mark. Did you forget to add it?`,
-    );
-  }
-  if (!insertion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain insertion mark. Did you forget to add it?`,
-    );
-  }
+  const { deletion, insertion } = getSuggestionMarks(node.type.schema);
 
   if (deletion.isInSet(node.marks)) {
     return null;
@@ -236,17 +221,7 @@ export function applySuggestions(
   state: EditorState,
   dispatch?: EditorView["dispatch"],
 ) {
-  const { deletion, insertion } = state.schema.marks;
-  if (!deletion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain deletion mark. Did you forget to add it?`,
-    );
-  }
-  if (!insertion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain insertion mark. Did you forget to add it?`,
-    );
-  }
+  const { deletion, insertion } = getSuggestionMarks(state.schema);
 
   const tr = state.tr;
   applySuggestionsToTransform(state.doc, tr, insertion, deletion);
@@ -265,17 +240,7 @@ export function applySuggestions(
  */
 export function applySuggestion(suggestionId: number): Command {
   return (state, dispatch) => {
-    const { deletion, insertion } = state.schema.marks;
-    if (!deletion) {
-      throw new Error(
-        `Failed to apply tracked changes to node: schema does not contain deletion mark. Did you forget to add it?`,
-      );
-    }
-    if (!insertion) {
-      throw new Error(
-        `Failed to apply tracked changes to node: schema does not contain insertion mark. Did you forget to add it?`,
-      );
-    }
+    const { deletion, insertion } = getSuggestionMarks(state.schema);
 
     const tr = state.tr;
     applySuggestionsToTransform(
@@ -304,17 +269,7 @@ export function revertSuggestions(
   state: EditorState,
   dispatch?: EditorView["dispatch"],
 ) {
-  const { deletion, insertion } = state.schema.marks;
-  if (!deletion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain deletion mark. Did you forget to add it?`,
-    );
-  }
-  if (!insertion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain insertion mark. Did you forget to add it?`,
-    );
-  }
+  const { deletion, insertion } = getSuggestionMarks(state.schema);
   const tr = state.tr;
   applySuggestionsToTransform(state.doc, tr, deletion, insertion);
   applyModificationsToTransform(tr.doc, tr, -1);
@@ -332,17 +287,7 @@ export function revertSuggestions(
  */
 export function revertSuggestion(suggestionId: number): Command {
   return (state, dispatch) => {
-    const { deletion, insertion } = state.schema.marks;
-    if (!deletion) {
-      throw new Error(
-        `Failed to apply tracked changes to node: schema does not contain deletion mark. Did you forget to add it?`,
-      );
-    }
-    if (!insertion) {
-      throw new Error(
-        `Failed to apply tracked changes to node: schema does not contain insertion mark. Did you forget to add it?`,
-      );
-    }
+    const { deletion, insertion } = getSuggestionMarks(state.schema);
 
     const tr = state.tr;
     applySuggestionsToTransform(
@@ -365,22 +310,9 @@ export function revertSuggestion(suggestionId: number): Command {
  */
 export function selectSuggestion(suggestionId: number): Command {
   return (state, dispatch) => {
-    const { deletion, insertion, modification } = state.schema.marks;
-    if (!deletion) {
-      throw new Error(
-        `Failed to apply tracked changes to node: schema does not contain deletion mark. Did you forget to add it?`,
-      );
-    }
-    if (!insertion) {
-      throw new Error(
-        `Failed to apply tracked changes to node: schema does not contain insertion mark. Did you forget to add it?`,
-      );
-    }
-    if (!modification) {
-      throw new Error(
-        `Failed to apply tracked changes to node: schema does not contain modification mark. Did you forget to add it?`,
-      );
-    }
+    const { deletion, insertion, modification } = getSuggestionMarks(
+      state.schema,
+    );
 
     let changeStart = null as number | null;
     let changeEnd = null as number | null;

src/decorations.ts

@@ -5,6 +5,7 @@ import {
   DecorationSet,
   type DecorationSource,
 } from "prosemirror-view";
+import { getSuggestionMarks } from "./utils.js";
 
 function pilcrow() {
   const span = document.createElement("span");
@@ -13,17 +14,7 @@ function pilcrow() {
 }
 
 export function getSuggestionDecorations(state: EditorState): DecorationSource {
-  const { deletion, insertion } = state.schema.marks;
-  if (!deletion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain deletion mark. Did you forget to add it?`,
-    );
-  }
-  if (!insertion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain insertion mark. Did you forget to add it?`,
-    );
-  }
+  const { deletion, insertion } = getSuggestionMarks(state.schema);
 
   const changeDecorations: Decoration[] = [];
   let lastParentNode: Node | null = null;

src/removeNodeMarkStep.ts

@@ -3,6 +3,7 @@ import { type EditorState, type Transaction } from "prosemirror-state";
 import { type RemoveNodeMarkStep, type Step } from "prosemirror-transform";
 
 import { rebasePos } from "./rebasePos.js";
+import { getSuggestionMarks } from "./utils.js";
 
 /**
  * Transform a remove node mark step into its equivalent tracked steps.
@@ -18,12 +19,7 @@ export function suggestRemoveNodeMarkStep(
   prevSteps: Step[],
   suggestionId: number,
 ) {
-  const { modification } = state.schema.marks;
-  if (!modification) {
-    throw new Error(
-      `Failed to apply modifications to node: schema does not contain modification mark. Did you forget to add it?`,
-    );
-  }
+  const { modification } = getSuggestionMarks(state.schema);
 
   const rebasedPos = rebasePos(step.pos, prevSteps, trackedTransaction.steps);
   const $pos = trackedTransaction.doc.resolve(rebasedPos);

src/replaceAroundStep.ts

@@ -16,6 +16,7 @@ import { applySuggestionsToSlice } from "./commands.js";
 import { rebasePos } from "./rebasePos.js";
 import { suggestRemoveNodeMarkStep } from "./removeNodeMarkStep.js";
 import { suggestReplaceStep } from "./replaceStep.js";
+import { getSuggestionMarks } from "./utils.js";
 
 /**
  * This detects and handles changes from `setNodeMarkup` so that these are tracked as a modification
@@ -37,12 +38,7 @@ function suggestSetNodeMarkup(
     step.gapFrom === step.from + 1 &&
     (step as ReplaceAroundStep & { structure: boolean }).structure
   ) {
-    const { modification } = state.schema.marks;
-    if (!modification) {
-      throw new Error(
-        `Failed to apply modifications to node: schema does not contain modification mark. Did you forget to add it?`,
-      );
-    }
+    const { modification } = getSuggestionMarks(state.schema);
 
     const newNode = step.slice.content.firstChild;
     let from = rebasePos(step.from, prevSteps, trackedTransaction.steps);

src/replaceStep.ts

@@ -8,6 +8,7 @@ import { type ReplaceStep, type Step } from "prosemirror-transform";
 
 import { findSuggestionMarkEnd } from "./findSuggestionMarkEnd.js";
 import { rebasePos } from "./rebasePos.js";
+import { getSuggestionMarks } from "./utils.js";
 
 /**
  * Transform a replace step into its equivalent tracked steps.
@@ -47,17 +48,7 @@ export function suggestReplaceStep(
   prevSteps: Step[],
   suggestionId: number,
 ) {
-  const { deletion, insertion } = state.schema.marks;
-  if (!deletion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain deletion mark. Did you forget to add it?`,
-    );
-  }
-  if (!insertion) {
-    throw new Error(
-      `Failed to apply tracked changes to node: schema does not contain insertion mark. Did you forget to add it?`,
-    );
-  }
+  const { deletion, insertion } = getSuggestionMarks(state.schema);
 
   // Check for insertion and deletion marks directly
   // adjacent to this step's boundaries. If they exist,

src/utils.ts

@@ -0,0 +1,35 @@
+import { type MarkType, type Schema } from "prosemirror-model";
+
+export interface SuggestionMarks {
+  insertion: MarkType;
+  deletion: MarkType;
+  modification: MarkType;
+}
+
+/**
+ * Get the suggestion mark types from a schema, with proper error handling.
+ * Throws an error if any of the required marks are not found.
+ */
+export function getSuggestionMarks(schema: Schema): SuggestionMarks {
+  const { insertion, deletion, modification } = schema.marks;
+
+  if (!insertion) {
+    throw new Error(
+      "Failed to find insertion mark in schema. Did you forget to add it?",
+    );
+  }
+
+  if (!deletion) {
+    throw new Error(
+      "Failed to find deletion mark in schema. Did you forget to add it?",
+    );
+  }
+
+  if (!modification) {
+    throw new Error(
+      "Failed to find modification mark in schema. Did you forget to add it?",
+    );
+  }
+
+  return { insertion, deletion, modification };
+}

src/withSuggestChanges.ts

@@ -20,6 +20,7 @@ import { suggestReplaceAroundStep } from "./replaceAroundStep.js";
 import { suggestReplaceStep } from "./replaceStep.js";
 import { type EditorView } from "prosemirror-view";
 import { isSuggestChangesEnabled, suggestChangesKey } from "./plugin.js";
+import { getSuggestionMarks } from "./utils.js";
 
 type StepHandler<S extends Step> = (
   trackedTransaction: Transaction,
@@ -99,22 +100,10 @@ export function transformToSuggestionTransaction(
   originalTransaction: Transaction,
   state: EditorState,
 ) {
-  const { deletion, insertion, modification } = state.schema.marks;
-  if (!deletion) {
-    throw new Error(
-      `Failed to transform to suggestion: schema does not contain deletion mark. Did you forget to add it?`,
-    );
-  }
-  if (!insertion) {
-    throw new Error(
-      `Failed to transform to suggestion: schema does not contain insertion mark. Did you forget to add it?`,
-    );
-  }
-  if (!modification) {
-    throw new Error(
-      `Failed to transform to suggestion: schema does not contain modification mark. Did you forget to add it?`,
-    );
-  }
+  // Validate that all required marks exist in the schema
+  const { deletion, insertion, modification } = getSuggestionMarks(
+    state.schema,
+  );
 
   // Find the highest change id in the document so far,
   // and use that as the starting point for new changes