Add visit(node, visitor) to 'yaml' (#225)

* Add visit(node, visitor) to 'yaml/util'
* Allow for a catch-all function visitor
* Include path as 2nd arg for visitor calls
* Skip visiting children by returning false
* End traversal when visit.BREAK is returned
* Use visit.SKIP rather than false to skip visiting children
* Return visit.REMOVE to remove the current node
* Throw if trying to remove a root node
* Add key argument to visitor calls
* Set next index in map/seq iteration by returning a number
* Add Alias visitor
* Move visit() export from 'yaml/util' to 'yaml'
* Update docs about YAML.visit()

README.md

@@ -44,6 +44,7 @@ const YAML = require('yaml')
   - [`#errors`](https://eemeli.org/yaml/#errors)
 - [`YAML.parseAllDocuments(str, options?): YAML.Document[]`](https://eemeli.org/yaml/#parsing-documents)
 - [`YAML.parseDocument(str, options?): YAML.Document`](https://eemeli.org/yaml/#parsing-documents)
+- [`YAML.visit(node, visitor)`](https://eemeli.org/yaml/#modifying-nodes)
 
 ```js
 import { Pair, YAMLMap, YAMLSeq } from 'yaml/types'

docs/01_intro.md

@@ -49,6 +49,7 @@ const YAML = require('yaml')
   - [`#errors`](#errors)
 - [`YAML.parseAllDocuments(str, options?): YAML.Document[]`](#parsing-documents)
 - [`YAML.parseDocument(str, options?): YAML.Document`](#parsing-documents)
+- [`YAML.visit(node, visitor)`](#modifying-nodes)
 
 ```js
 import { Pair, YAMLMap, YAMLSeq } from 'yaml/types'

docs/05_content_nodes.md

@@ -198,6 +198,71 @@ To construct a `YAMLSeq` or `YAMLMap`, use `doc.createNode()` with array, object
 
 Once created, normal array operations may be used to modify the `items` array. New `Pair` objects may created either by importing the class from `yaml/types` and using its `new Pair(key, value)` constructor, or by using the `doc.createPair(key, value, options?)` method. The latter will recursively wrap the `key` and `value` as nodes, and accepts the same options as `doc.createNode()`
 
+## Modifying Nodes
+
+```js
+const doc = YAML.parseDocument(`
+  - some values
+  - 42
+  - "3": a string
+    including: objects
+  - 1: a number
+`)
+
+const obs = doc.getIn([2, 'including'], true)
+obs.type = 'QUOTE_DOUBLE'
+
+YAML.visit(doc, {
+  Pair(_, pair) {
+    if (pair.key && pair.key.value === '3') return YAML.visit.REMOVE
+  },
+  Scalar(key, node) {
+    if (
+      key !== 'key' &&
+      typeof node.value === 'string' &&
+      node.type === 'PLAIN'
+    ) {
+      node.type = 'QUOTE_SINGLE'
+    }
+  }
+})
+
+String(doc)
+// - 'some values'
+// - 42
+// - including: "objects"
+// - 1: 'a number'
+```
+
+In general, it's safe to modify nodes manually, e.g. splicing the `items` array of a `YAMLMap` or changing its `type` from `'MAP'` to `'FLOW_MAP'`.
+For operations on nodes at a known location in the tree, it's probably easiest to use `doc.getIn(path, true)` to access them.
+For more complex or general operations, a visitor API is provided:
+
+#### `YAML.visit(node, visitor)`
+
+Apply a visitor to an AST node or document.
+
+Walks through the tree (depth-first) starting from `node`, calling a `visitor` function with three arguments:
+
+- `key`: For sequence values and map `Pair`, the node's index in the collection.
+  Within a `Pair`, `'key'` or `'value'`, correspondingly.
+  `null` for the root node.
+- `node`: The current node.
+- `path`: The ancestry of the current node.
+
+The return value of the visitor may be used to control the traversal:
+
+- `undefined` (default): Do nothing and continue
+- `YAML.visit.SKIP`: Do not visit the children of this node, continue with next sibling
+- `YAML.visit.BREAK`: Terminate traversal completely
+- `YAML.visit.REMOVE`: Remove the current node, then continue with the next one
+- `Node`: Replace the current node, then continue by visiting it
+- `number`: While iterating the items of a sequence or map, set the index of the next step.
+  This is useful especially if the index of the current node has changed.
+
+If `visitor` is a single function, it will be called with all values encountered in the tree, including e.g. `null` values.
+Alternatively, separate visitor functions may be defined for each `Map`, `Pair`, `Seq`, `Alias` and `Scalar` node.
+
 ## Comments
 
 ```js

index.d.ts

@@ -21,6 +21,65 @@ export interface ParsedCST extends Array<CST.Document> {
   setOrigRanges(): boolean
 }
 
+/**
+ * Apply a visitor to an AST node or document.
+ *
+ * Walks through the tree (depth-first) starting from `node`, calling a
+ * `visitor` function with three arguments:
+ *   - `key`: For sequence values and map `Pair`, the node's index in the
+ *     collection. Within a `Pair`, `'key'` or `'value'`, correspondingly.
+ *     `null` for the root node.
+ *   - `node`: The current node.
+ *   - `path`: The ancestry of the current node.
+ *
+ * The return value of the visitor may be used to control the traversal:
+ *   - `undefined` (default): Do nothing and continue
+ *   - `visit.SKIP`: Do not visit the children of this node, continue with next
+ *     sibling
+ *   - `visit.BREAK`: Terminate traversal completely
+ *   - `visit.REMOVE`: Remove the current node, then continue with the next one
+ *   - `Node`: Replace the current node, then continue by visiting it
+ *   - `number`: While iterating the items of a sequence or map, set the index
+ *     of the next step. This is useful especially if the index of the current
+ *     node has changed.
+ *
+ * If `visitor` is a single function, it will be called with all values
+ * encountered in the tree, including e.g. `null` values. Alternatively,
+ * separate visitor functions may be defined for each `Map`, `Pair`, `Seq`,
+ * `Alias` and `Scalar` node.
+ */
+export declare const visit: visit
+
+export type visitor<T> = (
+  key: number | 'key' | 'value' | null,
+  node: T,
+  path: Node[]
+) => void | symbol | number | Node
+
+export interface visit {
+  (
+    node: Node | Document,
+    visitor:
+      | visitor<any>
+      | {
+          Alias?: visitor<Alias>
+          Map?: visitor<YAMLMap>
+          Pair?: visitor<Pair>
+          Scalar?: visitor<Scalar>
+          Seq?: visitor<YAMLSeq>
+        }
+  ): void
+
+  /** Terminate visit traversal completely */
+  BREAK: symbol
+
+  /** Remove the current node */
+  REMOVE: symbol
+
+  /** Do not visit the children of the current node */
+  SKIP: symbol
+}
+
 /**
  * `yaml` defines document-specific options in three places: as an argument of
  * parse, create and stringify calls, in the values of `YAML.defaultOptions`,

src/ast/index.js

@@ -8,3 +8,4 @@ export { YAMLMap, findPair } from './YAMLMap.js'
 export { YAMLSeq } from './YAMLSeq.js'
 
 export { toJS } from './toJS.js'
+export { visit } from './visit.js'

src/ast/visit.js

@@ -0,0 +1,79 @@
+import { Type } from '../constants.js'
+import { Alias } from './Alias.js'
+import { Node } from './Node.js'
+import { Pair } from './Pair.js'
+import { Scalar } from './Scalar.js'
+import { YAMLMap } from './YAMLMap.js'
+import { YAMLSeq } from './YAMLSeq.js'
+
+const BREAK = Symbol('break visit')
+const SKIP = Symbol('skip children')
+const REMOVE = Symbol('remove node')
+
+function _visit(key, node, visitor, path) {
+  let ctrl = undefined
+  if (typeof visitor === 'function') ctrl = visitor(key, node, path)
+  else if (node instanceof YAMLMap) {
+    if (visitor.Map) ctrl = visitor.Map(key, node, path)
+  } else if (node instanceof YAMLSeq) {
+    if (visitor.Seq) ctrl = visitor.Seq(key, node, path)
+  } else if (node instanceof Pair) {
+    if (visitor.Pair) ctrl = visitor.Pair(key, node, path)
+  } else if (node instanceof Scalar) {
+    if (visitor.Scalar) ctrl = visitor.Scalar(key, node, path)
+  } else if (node instanceof Alias) {
+    if (visitor.Scalar) ctrl = visitor.Alias(key, node, path)
+  }
+
+  if (ctrl instanceof Node) {
+    const parent = path[path.length - 1]
+    if (parent instanceof YAMLMap || parent instanceof YAMLSeq) {
+      parent.items[key] = ctrl
+    } else if (parent instanceof Pair) {
+      if (key === 'key') parent.key = ctrl
+      else parent.value = ctrl
+    } else if (parent && parent.type === Type.DOCUMENT) {
+      parent.contents = ctrl
+    } else {
+      const pt = parent && parent.type
+      throw new Error(`Cannot replace node with ${pt} parent`)
+    }
+    return _visit(key, ctrl, visitor, path)
+  }
+
+  if (typeof ctrl !== 'symbol') {
+    if (node instanceof YAMLMap || node instanceof YAMLSeq) {
+      path = Object.freeze(path.concat(node))
+      for (let i = 0; i < node.items.length; ++i) {
+        const ci = _visit(i, node.items[i], visitor, path)
+        if (typeof ci === 'number') i = ci - 1
+        else if (ci === BREAK) return BREAK
+        else if (ci === REMOVE) {
+          node.items.splice(i, 1)
+          i -= 1
+        }
+      }
+    } else if (node instanceof Pair) {
+      path = Object.freeze(path.concat(node))
+      const ck = _visit('key', node.key, visitor, path)
+      if (ck === BREAK) return BREAK
+      else if (ck === REMOVE) node.key = null
+      const cv = _visit('value', node.value, visitor, path)
+      if (cv === BREAK) return BREAK
+      else if (cv === REMOVE) node.value = null
+    }
+  }
+
+  return ctrl
+}
+
+export function visit(node, visitor) {
+  if (node && node.type === Type.DOCUMENT) {
+    const cd = _visit(null, node.contents, visitor, Object.freeze([node]))
+    if (cd === REMOVE) node.contents = null
+  } else _visit(null, node, visitor, Object.freeze([]))
+}
+
+visit.BREAK = BREAK
+visit.SKIP = SKIP
+visit.REMOVE = REMOVE

src/index.js

@@ -4,6 +4,7 @@ import { Document } from './doc/Document.js'
 import { YAMLSemanticError } from './errors.js'
 import { warn } from './log.js'
 
+export { visit } from './ast/index.js'
 export { defaultOptions, scalarOptions } from './options.js'
 export { Document, parseCST }
 

tests/typings.ts

@@ -97,3 +97,20 @@ doc.contents = map
 
 const doc2 = new YAML.Document({ bizz: 'fuzz' })
 doc2.add(doc2.createPair('baz', 42))
+
+YAML.visit(doc, (key, node, path) => console.log(key, node, path))
+YAML.visit(doc, {
+  Scalar(key, node) {
+    if (key === 3) return 5
+    if (typeof node.value === 'number') return doc.createNode(node.value + 1)
+  },
+  Map(_, map) {
+    if (map.items.length > 3) return YAML.visit.SKIP
+  },
+  Pair(_, pair) {
+    if (pair.key.value === 'foo') return YAML.visit.REMOVE
+  },
+  Seq(_, seq) {
+    if (seq.items.length > 3) return YAML.visit.BREAK
+  }
+})