fix: Line comment handling in flow collections

src/nodes/Collection.ts

@@ -1,7 +1,6 @@
 import { createNode } from '../doc/createNode.js'
 import type { Schema } from '../schema/Schema.js'
-import { isCollection, isNode, isScalar, NodeBase, NODE_TYPE } from './Node.js'
-import type { Pair } from './Pair.js'
+import { isCollection, isPair, isScalar, NodeBase, NODE_TYPE } from './Node.js'
 
 export function collectionFromPath(
   schema: Schema,
@@ -143,8 +142,8 @@ export abstract class Collection extends NodeBase {
 
   hasAllNullValues(allowScalar?: boolean) {
     return this.items.every(node => {
-      if (!node || isNode(node)) return false
-      const n = (node as Pair).value
+      if (!isPair(node)) return false
+      const n = node.value
       return (
         n == null ||
         (allowScalar &&

src/stringify/stringifyCollection.ts

@@ -28,71 +28,76 @@ export function stringifyCollection(
   const inFlow = flow || ctx.inFlow
   if (inFlow) itemIndent += indentStep
   ctx = Object.assign({}, ctx, { indent: itemIndent, inFlow, type: null })
+
+  let singleLineOutput = true
   let chompKeep = false // flag for the preceding node's status
-  let hasItemWithNewLine = false
   const nodes = items.reduce((nodes: StringifyNode[], item, i) => {
     let comment: string | null = null
     if (isNode(item)) {
-      if (!chompKeep && item.spaceBefore) {
-        nodes.push({ comment: true, str: '' })
-        if (inFlow) hasItemWithNewLine = true
-      }
+      if (!chompKeep && item.spaceBefore) nodes.push({ comment: true, str: '' })
       if (item.commentBefore) {
         // This match will always succeed on a non-empty string
         for (const line of item.commentBefore.match(/^.*$/gm) as string[])
           nodes.push({ comment: true, str: `#${line}` })
-        if (inFlow) hasItemWithNewLine = true
       }
       if (item.comment) {
         comment = item.comment
-        if (inFlow) hasItemWithNewLine = true
+        singleLineOutput = false
       }
     } else if (isPair(item)) {
-      if (isNode(item.key)) {
-        if (!chompKeep && item.key.spaceBefore) {
-          nodes.push({ comment: true, str: '' })
-          if (inFlow) hasItemWithNewLine = true
-        }
-        if (item.key.commentBefore) {
+      const ik = isNode(item.key) ? item.key : null
+      if (ik) {
+        if (!chompKeep && ik.spaceBefore) nodes.push({ comment: true, str: '' })
+        if (ik.commentBefore) {
           // This match will always succeed on a non-empty string
-          for (const line of item.key.commentBefore.match(/^.*$/gm) as string[])
+          for (const line of ik.commentBefore.match(/^.*$/gm) as string[])
             nodes.push({ comment: true, str: `#${line}` })
-          if (inFlow) hasItemWithNewLine = true
         }
-        if (inFlow && item.key.comment) hasItemWithNewLine = true
+        if (ik.comment) singleLineOutput = false
+      }
+
+      if (inFlow) {
+        const iv = isNode(item.value) ? item.value : null
+        if (iv) {
+          if (iv.comment) comment = iv.comment
+          if (iv.comment || iv.commentBefore) singleLineOutput = false
+        } else if (item.value == null && ik && ik.comment) {
+          comment = ik.comment
+        }
       }
-      if (
-        inFlow &&
-        isNode(item.value) &&
-        (item.value.commentBefore || item.value.comment)
-      )
-        hasItemWithNewLine = true
     }
+
     chompKeep = false
     let str = stringify(
       item,
       ctx,
       () => (comment = null),
       () => (chompKeep = true)
     )
-    if (inFlow && !hasItemWithNewLine && str.includes('\n'))
-      hasItemWithNewLine = true
     if (inFlow && i < items.length - 1) str += ','
     str = addComment(str, itemIndent, comment)
     if (chompKeep && (comment || inFlow)) chompKeep = false
     nodes.push({ comment: false, str })
     return nodes
   }, [])
+
   let str: string
   if (nodes.length === 0) {
     str = flowChars.start + flowChars.end
   } else if (inFlow) {
     const { start, end } = flowChars
     const strings = nodes.map(n => n.str)
+    let singleLineLength = 2
+    for (const node of nodes) {
+      if (node.comment || node.str.includes('\n')) {
+        singleLineOutput = false
+        break
+      }
+      singleLineLength += node.str.length + 2
+    }
     if (
-      hasItemWithNewLine ||
-      strings.reduce((sum, str) => sum + str.length + 2, 2) >
-        Collection.maxFlowStringSingleLineLength
+      !singleLineOutput ||
+      singleLineLength > Collection.maxFlowStringSingleLineLength
     ) {
       str = start
       for (const s of strings) {

src/stringify/stringifyPair.ts

@@ -7,7 +7,7 @@ import { stringify, StringifyContext } from './stringify.js'
 export function stringifyPair(
   { key, value }: Readonly<Pair>,
   ctx: StringifyContext,
-  _onComment?: () => void,
+  onComment?: () => void,
   onChompKeep?: () => void
 ) {
   const {
@@ -41,11 +41,12 @@ export function stringifyPair(
     implicitKey: !explicitKey && (simpleKeys || !allNullValues),
     indent: indent + indentStep
   })
+  let keyCommentDone = false
   let chompKeep = false
   let str = stringify(
     key,
     ctx,
-    () => (keyComment = null),
+    () => (keyCommentDone = true),
     () => (chompKeep = true)
   )
 
@@ -57,15 +58,18 @@ export function stringifyPair(
     explicitKey = true
   }
 
-  if (
-    (allNullValues && (!simpleKeys || ctx.inFlow)) ||
-    (value == null && (explicitKey || ctx.inFlow))
-  ) {
-    str = addComment(str, ctx.indent, keyComment)
+  if (ctx.inFlow) {
+    if (allNullValues || value == null) {
+      if (keyCommentDone && onComment) onComment()
+      return explicitKey ? `? ${str}` : str
+    }
+  } else if ((allNullValues && !simpleKeys) || (value == null && explicitKey)) {
+    if (keyCommentDone) keyComment = null
     if (chompKeep && !keyComment && onChompKeep) onChompKeep()
-    return ctx.inFlow && !explicitKey ? str : `? ${str}`
+    return addComment(`? ${str}`, ctx.indent, keyComment)
   }
 
+  if (keyCommentDone) keyComment = null
   str = explicitKey
     ? `? ${addComment(str, ctx.indent, keyComment)}\n${indent}:`
     : addComment(`${str}:`, ctx.indent, keyComment)
@@ -99,10 +103,12 @@ export function stringifyPair(
     // If indentSeq === false, consider '- ' as part of indentation where possible
     ctx.indent = ctx.indent.substr(2)
   }
+
+  let valueCommentDone = false
   const valueStr = stringify(
     value,
     ctx,
-    () => (valueComment = null),
+    () => (valueCommentDone = true),
     () => (chompKeep = true)
   )
   let ws = ' '
@@ -112,6 +118,13 @@ export function stringifyPair(
     const flow = valueStr[0] === '[' || valueStr[0] === '{'
     if (!flow || valueStr.includes('\n')) ws = `\n${ctx.indent}`
   } else if (valueStr[0] === '\n') ws = ''
-  if (chompKeep && !valueComment && onChompKeep) onChompKeep()
-  return addComment(str + ws + valueStr, ctx.indent, valueComment)
+
+  if (ctx.inFlow) {
+    if (valueCommentDone && onComment) onComment()
+    return str + ws + valueStr
+  } else {
+    if (valueCommentDone) valueComment = null
+    if (chompKeep && !valueComment && onChompKeep) onChompKeep()
+    return addComment(str + ws + valueStr, ctx.indent, valueComment)
+  }
 }

tests/doc/comments.js

@@ -471,7 +471,7 @@ describe('stringify comments', () => {
     })
   })
 
-  describe('flow collection commens', () => {
+  describe('flow collection comments', () => {
     test('line comment after , in seq', () => {
       const doc = YAML.parseDocument(source`
         [ a, #c0
@@ -485,17 +485,17 @@ describe('stringify comments', () => {
       `)
     })
 
-    test.skip('line comment after , in map', () => {
+    test('line comment after , in map', () => {
       const doc = YAML.parseDocument(source`
         { a, #c0
           b: c, #c1
           d #c2
         }`)
       expect(String(doc)).toBe(source`
         {
-          ? a, #c0
+          a, #c0
           b: c, #c1
-          ? d #c2
+          d #c2
         }
       `)
     })