Preserve indentation when using parseDocument

src/compose/resolve-block-map.ts

@@ -102,7 +102,24 @@ export function resolveBlockMap(
         : composeEmptyNode(ctx, offset, sep, null, valueProps, onError)
       offset = valueNode.range[2]
       const pair = new Pair(keyNode, valueNode)
-      if (ctx.options.keepSourceTokens) pair.srcToken = collItem
+      if (ctx.options.keepSourceTokens) {
+        pair.srcToken = collItem
+
+        // Check to see if the key tokens exist and have an indentation set.
+        // If so, we can compute the difference between them and preserve it
+        // if the preserveCollectionIndentation option is set.
+        const keyIndent: number | undefined = (collItem?.key as any)?.indent
+        const valueIndent: number | undefined = (collItem?.value as any)?.indent
+
+        if (
+          ctx.options.preserveCollectionIndentation &&
+          keyIndent !== undefined &&
+          valueIndent !== undefined &&
+          valueIndent >= keyIndent
+        ) {
+          pair.srcIndentStep = valueIndent - keyIndent
+        }
+      }
       map.items.push(pair)
     } else {
       // key with no value

src/nodes/Pair.ts

@@ -29,6 +29,12 @@ export class Pair<K = unknown, V = unknown> {
   /** The CST token that was composed into this pair.  */
   declare srcToken?: CollectionItem
 
+  /**
+   * The indentation step between key and value in the source, to be
+   * preserved during stringification.
+   */
+  declare srcIndentStep?: number
+
   constructor(key: K, value: V | null = null) {
     Object.defineProperty(this, NODE_TYPE, { value: PAIR })
     this.key = key

src/options.ts

@@ -27,6 +27,17 @@ export type ParseOptions = {
    */
   keepSourceTokens?: boolean
 
+  /**
+   * When parsing a document, stores indentation levels on each collection node,
+   * allowing them to be preserved when latery stringified.
+   *
+   * If 'keepSourceTokens' is not `true`, this has no affect, since source tokens
+   * must be provided in order to determine indentation.
+   *
+   * Default: `false`
+   */
+  preserveCollectionIndentation?: boolean
+
   /**
    * If set, newlines will be tracked, to allow for `lineCounter.linePos(offset)`
    * to provide the `{ line, col }` positions within the input.

src/stringify/stringifyPair.ts

@@ -5,18 +5,39 @@ import { stringify, StringifyContext } from './stringify.js'
 import { addComment, stringifyComment } from './stringifyComment.js'
 
 export function stringifyPair(
-  { key, value }: Readonly<Pair>,
+  { key, value, srcIndentStep }: Readonly<Pair>,
   ctx: StringifyContext,
   onComment?: () => void,
   onChompKeep?: () => void
 ) {
-  const {
+  let {
     allNullValues,
     doc,
     indent,
     indentStep,
     options: { indentSeq, simpleKeys }
   } = ctx
+  if (srcIndentStep !== undefined) {
+    // If the pair originally had some indentation step, preserve that
+    // value.
+    if (srcIndentStep > 0 || (srcIndentStep === 0 && isSeq(value))) {
+      // Indentation can only be preserved if it's positive, or if it's 0
+      // and the item to render is a seq, since:
+
+      // foo:
+      // - a
+      // - b
+      // - c
+
+      // is a valid seq with 0 indentStep.
+
+      // Note that ctx.indentStep is *not* modified, so later indentations
+      // will still use the original indentStep if not being preserved from
+      // input.
+      indentStep = ' '.repeat(srcIndentStep)
+    }
+  }
+
   let keyComment = (isNode(key) && key.comment) || null
   if (simpleKeys) {
     if (keyComment) {

tests/preserve-indentation.ts

@@ -0,0 +1,160 @@
+import { parseDocument } from 'yaml'
+
+describe('preserveCollectionIndentation', () => {
+  // This sample document has very unusual indentation, which helps
+  // to ensure that it really does end up being preserved exactly.
+  const sample = `
+a:
+      b:
+       c: d
+       more: e
+
+big_indent:
+         - a
+         - b
+         - c
+more:
+   # A comment goes here,
+   # which spans multiple lines.
+
+   # Ideally, a blank line is still preserved as well,
+   # so that large comment blocks can be separated.
+   a:
+    - x # after the item
+    # between items
+    - y:
+       further:
+        indentation:
+        - is
+        - possible
+        - even:
+           tho:
+                  it:
+                         changes: a lot
+    - z
+    # after last
+`
+
+  test('preserveCollectionIndentation: toString() preserve document indentation', () => {
+    const document = parseDocument(sample, {
+      keepSourceTokens: true,
+      preserveCollectionIndentation: true
+    })
+    const roundtrippedSource = document.toString()
+
+    expect(roundtrippedSource.trim()).toEqual(sample.trim())
+  })
+
+  // sample2 corresponds to the JSON `{"foo": ["a", "b", "c"]}`
+  // The Seq is not indented at all, and we can preserve that when
+  // parsing.
+  const sample2 = `
+foo:
+- a
+- b
+- c
+`
+  // However, when we edit the field to replace the list with an object,
+  // some indentation is now required. In this case, we default to the
+  // original 'indentStep' value (e.g. 2 in this case) because some
+  // level of indentation is required.
+  const sample3 = `
+foo:
+  a: 1
+  b: 2
+  c: 3
+`
+
+  test('preserveCollectionIndentation: produces correct yaml when preserving indentation with edits', () => {
+    const document = parseDocument(sample2, {
+      keepSourceTokens: true,
+      preserveCollectionIndentation: true
+    })
+    expect(document.toString().trim()).toEqual(sample2.trim())
+
+    // When replacing the item, we now need indentation, since otherwise
+    // the document has a different structure than it initially did.
+    document.set('foo', { a: 1, b: 2, c: 3 })
+    expect(document.toString().trim()).toEqual(sample3.trim())
+  })
+
+  const combined = `
+a:
+  b:
+    c: d
+    more: e
+
+big_indent:
+  - a
+  - b
+  - c
+more:
+  # A comment goes here,
+  # which spans multiple lines.
+
+  # Ideally, a blank line is still preserved as well,
+  # so that large comment blocks can be separated.
+  a:
+    - x # after the item
+    # between items
+    - y:
+        further:
+          indentation:
+            - is
+            - possible
+            - even:
+                tho:
+                  it:
+                    changes: a lot
+    - z
+    # after last
+preservedDocument:
+  a:
+        b:
+         c: d
+         more: e
+
+  big_indent:
+           - a
+           - b
+           - c
+  more:
+     # A comment goes here,
+     # which spans multiple lines.
+
+     # Ideally, a blank line is still preserved as well,
+     # so that large comment blocks can be separated.
+     a:
+      - x # after the item
+      # between items
+      - y:
+         further:
+          indentation:
+          - is
+          - possible
+          - even:
+             tho:
+                    it:
+                           changes: a lot
+      - z
+      # after last
+`
+
+  test('preserveCollectionIndentation: documents with preserved indentation can be inserted into other documents', () => {
+    const document = parseDocument(sample, {
+      keepSourceTokens: true,
+      preserveCollectionIndentation: true
+    })
+
+    // In this new document, we purposefully skip preserving indentation:
+    const newDocument = parseDocument(sample, {
+      preserveCollectionIndentation: false
+    })
+
+    // The indentation-preserved document is inserted into the "main" document,
+    // and maintains its original indentation level.
+    newDocument.set('preservedDocument', document)
+
+    expect(newDocument.toString().trim()).toEqual(combined.trim())
+  })
+})

tests/yaml-test-suite.ts

@@ -98,6 +98,21 @@ for (const dir of testDirs) {
             const docs2 = parseAllDocuments(src2, { resolveKnownTags: false })
             testJsonMatch(docs2, json)
           })
+
+          test('stringify+re-parse when preserving indentation', () => {
+            const roundTripDocuments = parseAllDocuments(yaml, {
+              resolveKnownTags: false,
+              keepSourceTokens: true,
+              preserveCollectionIndentation: true
+            })
+            testJsonMatch(roundTripDocuments, json)
+
+            const src2 =
+              docs.map(doc => String(doc).replace(/\n$/, '')).join('\n...\n') +
+              '\n'
+            const docs2 = parseAllDocuments(src2, { resolveKnownTags: false })
+            testJsonMatch(docs2, json)
+          })
         }
 
         const outYaml = load('out.yaml')