Indent-based folds for YAML editor

src/components/ha-code-editor.ts

@@ -123,10 +123,14 @@ export class HaCodeEditor extends ReactiveElement {
     }
     const transactions: TransactionSpec[] = [];
     if (changedProps.has("mode")) {
+      // TODO: not sure how to handle things here
       transactions.push({
-        effects: this._loadedCodeMirror!.langCompartment!.reconfigure(
-          this._mode
-        ),
+        effects: [
+          this._loadedCodeMirror!.langCompartment!.reconfigure(this._mode),
+          this._loadedCodeMirror!.foldingCompartment.reconfigure(
+            this._getFoldingExtensions()
+          ),
+        ],
       });
     }
     if (changedProps.has("readOnly")) {
@@ -194,6 +198,9 @@ export class HaCodeEditor extends ReactiveElement {
         this.linewrap ? this._loadedCodeMirror.EditorView.lineWrapping : []
       ),
       this._loadedCodeMirror.EditorView.updateListener.of(this._onUpdate),
+      this._loadedCodeMirror.foldingCompartment.of(
+        this._getFoldingExtensions()
+      ),
     ];
 
     if (!this.readOnly) {
@@ -311,6 +318,17 @@ export class HaCodeEditor extends ReactiveElement {
     fireEvent(this, "value-changed", { value: this._value });
   };
 
+  private _getFoldingExtensions = (): Extension => {
+    if (this.mode === "yaml") {
+      return [
+        this._loadedCodeMirror!.foldGutter(),
+        this._loadedCodeMirror!.foldingOnIndent,
+      ];
+    }
+
+    return [];
+  };
+
   static get styles(): CSSResultGroup {
     return css`
       :host(.error-state) .cm-gutters {

src/resources/codemirror.ts

@@ -1,5 +1,6 @@
 import { indentLess, indentMore } from "@codemirror/commands";
 import {
+  foldService,
   HighlightStyle,
   StreamLanguage,
   syntaxHighlighting,
@@ -12,7 +13,7 @@ import { tags } from "@lezer/highlight";
 
 export { autocompletion } from "@codemirror/autocomplete";
 export { defaultKeymap, history, historyKeymap } from "@codemirror/commands";
-export { highlightingFor } from "@codemirror/language";
+export { highlightingFor, foldGutter } from "@codemirror/language";
 export { highlightSelectionMatches, searchKeymap } from "@codemirror/search";
 export { EditorState } from "@codemirror/state";
 export {
@@ -34,6 +35,7 @@ export const langs = {
 export const langCompartment = new Compartment();
 export const readonlyCompartment = new Compartment();
 export const linewrapCompartment = new Compartment();
+export const foldingCompartment = new Compartment();
 
 export const tabKeyBindings: KeyBinding[] = [
   { key: "Tab", run: indentMore },
@@ -270,3 +272,42 @@ const haHighlightStyle = HighlightStyle.define([
 ]);
 
 export const haSyntaxHighlighting = syntaxHighlighting(haHighlightStyle);
+
+// A folding service for indent-based languages such as YAML.
+export const foldingOnIndent = foldService.of((state, from, to) => {
+  const line = state.doc.lineAt(from);
+  const lineCount = state.doc.lines;
+  const indent = line.text.search(/\S|$/); // Indent level of the first line
+  let foldStart = from; // Start of the fold
+  let foldEnd = to; // End of the fold
+
+  // Check if the next line is on a deeper indent level
+  // If so, continue subsequent lines
+  // If not, go on with the foldEnd
+  let nextLine = line;
+  while (nextLine.number < lineCount) {
+    nextLine = state.doc.line(nextLine.number + 1); // Next line
+    const nextIndent = nextLine.text.search(/\S|$/); // Indent level of the next line
+
+    // If the next line is on a deeper indent level, add it to the fold
+    if (nextIndent > indent) {
+      // include this line in the fold and continue
+      foldEnd = nextLine.to;
+    } else {
+      // If the next line is not on a deeper indent level, we found the end of the region
+      break;
+    }
+  }
+
+  // Don't create fold if it's a single line
+  if (state.doc.lineAt(foldStart).number === state.doc.lineAt(foldEnd).number) {
+    return null;
+  }
+
+  // Set the fold start to the end of the first line
+  // With this, the fold will not include the first line
+  foldStart = line.to;
+
+  // Return a fold that covers the entire indent level
+  return { from: foldStart, to: foldEnd };
+});