feat: Improve/specify/fix TS types for Document & collection access m…

…ethods (#383)

* feat: generic argument support for node methods

* fix: Types of collection access methods

* feat: Specify Document content & createNode() types more precisely

* feat: Export types for YAMLOmap and YAMLSet

* test: Port collection-access tests to TypeScript

* test: Updated tests and types for collection access

* style: Reversed generic parameter rename

* fix: Corrected overload signatures

* test: Corrected test assertion style

Co-authored-by: Eemeli Aro <eemeli@gmail.com>

src/doc/Document.ts

@@ -7,6 +7,7 @@ import {
   isNode,
   isScalar,
   Node,
+  NodeType,
   NODE_TYPE,
   ParsedNode,
   Range
@@ -41,7 +42,7 @@ export declare namespace Document {
   }
 }
 
-export class Document<T = unknown> {
+export class Document<T extends Node = Node> {
   readonly [NODE_TYPE]: symbol
 
   /** A comment before this Document */
@@ -192,12 +193,12 @@ export class Document<T = unknown> {
    * Convert any value into a `Node` using the current schema, recursively
    * turning objects into collections.
    */
-  createNode(value: unknown, options?: CreateNodeOptions): Node
-  createNode(
-    value: unknown,
+  createNode<T = unknown>(value: T, options?: CreateNodeOptions): NodeType<T>
+  createNode<T = unknown>(
+    value: T,
     replacer: Replacer | CreateNodeOptions | null,
     options?: CreateNodeOptions
-  ): Node
+  ): NodeType<T>
   createNode(
     value: unknown,
     replacer?: Replacer | CreateNodeOptions | null,
@@ -264,15 +265,15 @@ export class Document<T = unknown> {
    * Removes a value from the document.
    * @returns `true` if the item was found and removed.
    */
-  delete(key: any) {
+  delete(key: unknown): boolean {
     return assertCollection(this.contents) ? this.contents.delete(key) : false
   }
 
   /**
    * Removes a value from the document.
    * @returns `true` if the item was found and removed.
    */
-  deleteIn(path: Iterable<unknown>) {
+  deleteIn(path: Iterable<unknown> | null): boolean {
     if (isEmptyPath(path)) {
       if (this.contents == null) return false
       this.contents = null
@@ -288,7 +289,7 @@ export class Document<T = unknown> {
    * scalar values from their surrounding node; to disable set `keepScalar` to
    * `true` (collections are always returned intact).
    */
-  get(key: unknown, keepScalar?: boolean) {
+  get(key: unknown, keepScalar?: boolean): unknown {
     return isCollection(this.contents)
       ? this.contents.get(key, keepScalar)
       : undefined
@@ -299,7 +300,7 @@ export class Document<T = unknown> {
    * scalar values from their surrounding node; to disable set `keepScalar` to
    * `true` (collections are always returned intact).
    */
-  getIn(path: Iterable<unknown>, keepScalar?: boolean) {
+  getIn(path: Iterable<unknown> | null, keepScalar?: boolean): unknown {
     if (isEmptyPath(path))
       return !keepScalar && isScalar(this.contents)
         ? this.contents.value
@@ -312,14 +313,14 @@ export class Document<T = unknown> {
   /**
    * Checks if the document includes a value with the key `key`.
    */
-  has(key: unknown) {
+  has(key: unknown): boolean {
     return isCollection(this.contents) ? this.contents.has(key) : false
   }
 
   /**
    * Checks if the document includes a value at `path`.
    */
-  hasIn(path: Iterable<unknown>) {
+  hasIn(path: Iterable<unknown> | null): boolean {
     if (isEmptyPath(path)) return this.contents !== undefined
     return isCollection(this.contents) ? this.contents.hasIn(path) : false
   }
@@ -328,7 +329,7 @@ export class Document<T = unknown> {
    * Sets a value in this document. For `!!set`, `value` needs to be a
    * boolean to add/remove the item from the set.
    */
-  set(key: any, value: unknown) {
+  set(key: any, value: unknown): void {
     if (this.contents == null) {
       this.contents = collectionFromPath(
         this.schema,
@@ -344,7 +345,7 @@ export class Document<T = unknown> {
    * Sets a value in this document. For `!!set`, `value` needs to be a
    * boolean to add/remove the item from the set.
    */
-  setIn(path: Iterable<unknown>, value: unknown) {
+  setIn(path: Iterable<unknown> | null, value: unknown): void {
     if (isEmptyPath(path)) this.contents = value as T
     else if (this.contents == null) {
       this.contents = collectionFromPath(

src/index.ts

@@ -48,6 +48,8 @@ export {
 
 export type { TagId, Tags } from './schema/tags'
 export type { CollectionTag, ScalarTag } from './schema/types'
+export type { YAMLOMap } from './schema/yaml-1.1/omap'
+export type { YAMLSet } from './schema/yaml-1.1/set'
 
 export {
   asyncVisitor,

src/nodes/Collection.ts

@@ -36,8 +36,11 @@ export function collectionFromPath(
   })
 }
 
-// null, undefined, or an empty non-string iterable (e.g. [])
-export const isEmptyPath = (path: Iterable<unknown> | null | undefined) =>
+// Type guard is intentionally a little wrong so as to be more useful,
+// as it does not cover untypable empty non-string iterables (e.g. []).
+export const isEmptyPath = (
+  path: Iterable<unknown> | null | undefined
+): path is null | undefined =>
   path == null ||
   (typeof path === 'object' && !!path[Symbol.iterator]().next().done)
 

src/nodes/Node.ts

@@ -7,7 +7,20 @@ import type { Scalar } from './Scalar.js'
 import type { YAMLMap } from './YAMLMap.js'
 import type { YAMLSeq } from './YAMLSeq.js'
 
-export type Node = Alias | Scalar | YAMLMap | YAMLSeq
+export type Node<T = unknown> =
+  | Alias
+  | Scalar<T>
+  | YAMLMap<unknown, T>
+  | YAMLSeq<T>
+
+/** Utility type mapper */
+export type NodeType<T> = T extends string | number | bigint | boolean | null
+  ? Scalar<T>
+  : T extends Array<any>
+  ? YAMLSeq<NodeType<T[number]>>
+  : T extends { [key: string | number]: any }
+  ? YAMLMap<NodeType<keyof T>, NodeType<T[keyof T]>>
+  : Node
 
 export type ParsedNode =
   | Alias.Parsed
@@ -28,22 +41,30 @@ export const NODE_TYPE = Symbol.for('yaml.node.type')
 export const isAlias = (node: any): node is Alias =>
   !!node && typeof node === 'object' && node[NODE_TYPE] === ALIAS
 
-export const isDocument = (node: any): node is Document =>
+export const isDocument = <T extends Node = Node>(
+  node: any
+): node is Document<T> =>
   !!node && typeof node === 'object' && node[NODE_TYPE] === DOC
 
-export const isMap = (node: any): node is YAMLMap =>
+export const isMap = <K = unknown, V = unknown>(
+  node: any
+): node is YAMLMap<K, V> =>
   !!node && typeof node === 'object' && node[NODE_TYPE] === MAP
 
-export const isPair = (node: any): node is Pair =>
+export const isPair = <K = unknown, V = unknown>(
+  node: any
+): node is Pair<K, V> =>
   !!node && typeof node === 'object' && node[NODE_TYPE] === PAIR
 
-export const isScalar = (node: any): node is Scalar =>
+export const isScalar = <T = unknown>(node: any): node is Scalar<T> =>
   !!node && typeof node === 'object' && node[NODE_TYPE] === SCALAR
 
-export const isSeq = (node: any): node is YAMLSeq =>
+export const isSeq = <T = unknown>(node: any): node is YAMLSeq<T> =>
   !!node && typeof node === 'object' && node[NODE_TYPE] === SEQ
 
-export function isCollection(node: any): node is YAMLMap | YAMLSeq {
+export function isCollection<K = unknown, V = unknown>(
+  node: any
+): node is YAMLMap<K, V> | YAMLSeq<V> {
   if (node && typeof node === 'object')
     switch (node[NODE_TYPE]) {
       case MAP:
@@ -53,7 +74,7 @@ export function isCollection(node: any): node is YAMLMap | YAMLSeq {
   return false
 }
 
-export function isNode(node: any): node is Node {
+export function isNode<T = unknown>(node: any): node is Node<T> {
   if (node && typeof node === 'object')
     switch (node[NODE_TYPE]) {
       case ALIAS:
@@ -65,7 +86,9 @@ export function isNode(node: any): node is Node {
   return false
 }
 
-export const hasAnchor = (node: unknown): node is Scalar | YAMLMap | YAMLSeq =>
+export const hasAnchor = <K = unknown, V = unknown>(
+  node: unknown
+): node is Scalar<V> | YAMLMap<K, V> | YAMLSeq<V> =>
   (isScalar(node) || isCollection(node)) && !!node.anchor
 
 export abstract class NodeBase {

src/nodes/YAMLMap.ts

@@ -6,7 +6,7 @@ import { addPairToJSMap } from './addPairToJSMap.js'
 import { Collection } from './Collection.js'
 import { isPair, isScalar, MAP, ParsedNode, Range } from './Node.js'
 import { Pair } from './Pair.js'
-import { isScalarValue } from './Scalar.js'
+import { isScalarValue, Scalar } from './Scalar.js'
 import type { ToJSContext } from './toJS.js'
 
 export function findPair<K = unknown, V = unknown>(
@@ -51,12 +51,12 @@ export class YAMLMap<K = unknown, V = unknown> extends Collection {
    * @param overwrite - If not set `true`, using a key that is already in the
    *   collection will throw. Otherwise, overwrites the previous value.
    */
-  add(pair: Pair<K, V> | { key: K; value: V }, overwrite?: boolean) {
+  add(pair: Pair<K, V> | { key: K; value: V }, overwrite?: boolean): void {
     let _pair: Pair<K, V>
     if (isPair(pair)) _pair = pair
     else if (!pair || typeof pair !== 'object' || !('key' in pair)) {
       // In TypeScript, this never happens.
-      _pair = new Pair<K, V>(pair as any, (pair as any).value)
+      _pair = new Pair<K, V>(pair as any, (pair as any)?.value)
     } else _pair = new Pair(pair.key, pair.value)
 
     const prev = findPair(this.items, _pair.key)
@@ -76,24 +76,27 @@ export class YAMLMap<K = unknown, V = unknown> extends Collection {
     }
   }
 
-  delete(key: K) {
+  delete(key: unknown): boolean {
     const it = findPair(this.items, key)
     if (!it) return false
     const del = this.items.splice(this.items.indexOf(it), 1)
     return del.length > 0
   }
 
-  get(key: K, keepScalar?: boolean) {
+  get(key: unknown, keepScalar: true): Scalar<V> | undefined
+  get(key: unknown, keepScalar?: false): V | undefined
+  get(key: unknown, keepScalar?: boolean): V | Scalar<V> | undefined
+  get(key: unknown, keepScalar?: boolean): V | Scalar<V> | undefined {
     const it = findPair(this.items, key)
     const node = it?.value
-    return !keepScalar && isScalar(node) ? node.value : node
+    return (!keepScalar && isScalar<V>(node) ? node.value : node) ?? undefined
   }
 
-  has(key: K) {
+  has(key: unknown): boolean {
     return !!findPair(this.items, key)
   }
 
-  set(key: K, value: V) {
+  set(key: K, value: V): void {
     this.add(new Pair(key, value), true)
   }
 

src/nodes/YAMLSeq.ts

@@ -5,7 +5,7 @@ import { stringifyCollection } from '../stringify/stringifyCollection.js'
 import { Collection } from './Collection.js'
 import { isScalar, ParsedNode, Range, SEQ } from './Node.js'
 import type { Pair } from './Pair.js'
-import { isScalarValue } from './Scalar.js'
+import { isScalarValue, Scalar } from './Scalar.js'
 import { toJS, ToJSContext } from './toJS.js'
 
 export declare namespace YAMLSeq {
@@ -29,7 +29,7 @@ export class YAMLSeq<T = unknown> extends Collection {
     super(SEQ, schema)
   }
 
-  add(value: T) {
+  add(value: T): void {
     this.items.push(value)
   }
 
@@ -41,7 +41,7 @@ export class YAMLSeq<T = unknown> extends Collection {
    *
    * @returns `true` if the item was found and removed.
    */
-  delete(key: unknown) {
+  delete(key: unknown): boolean {
     const idx = asItemIndex(key)
     if (typeof idx !== 'number') return false
     const del = this.items.splice(idx, 1)
@@ -56,11 +56,20 @@ export class YAMLSeq<T = unknown> extends Collection {
    * `key` must contain a representation of an integer for this to succeed.
    * It may be wrapped in a `Scalar`.
    */
-  get(key: unknown, keepScalar?: boolean) {
+  get(key: unknown, keepScalar: true): Scalar<T> | undefined
+  get(key: unknown, keepScalar?: false): T | undefined
+  get(
+    key: unknown,
+    keepScalar?: boolean
+  ): T | Scalar<T> | undefined
+  get(
+    key: unknown,
+    keepScalar?: boolean
+  ): T | Scalar<T> | undefined {
     const idx = asItemIndex(key)
     if (typeof idx !== 'number') return undefined
     const it = this.items[idx]
-    return !keepScalar && isScalar(it) ? it.value : it
+    return !keepScalar && isScalar<T>(it) ? it.value : it
   }
 
   /**
@@ -69,7 +78,7 @@ export class YAMLSeq<T = unknown> extends Collection {
    * `key` must contain a representation of an integer for this to succeed.
    * It may be wrapped in a `Scalar`.
    */
-  has(key: unknown) {
+  has(key: unknown): boolean {
     const idx = asItemIndex(key)
     return typeof idx === 'number' && idx < this.items.length
   }
@@ -81,7 +90,7 @@ export class YAMLSeq<T = unknown> extends Collection {
    * If `key` does not contain a representation of an integer, this will throw.
    * It may be wrapped in a `Scalar`.
    */
-  set(key: unknown, value: T) {
+  set(key: unknown, value: T): void {
     const idx = asItemIndex(key)
     if (typeof idx !== 'number')
       throw new Error(`Expected a valid index, not ${key}.`)

src/schema/yaml-1.1/set.ts

@@ -35,7 +35,11 @@ export class YAMLSet<T = unknown> extends YAMLMap<T, Scalar<null> | null> {
     if (!prev) this.items.push(pair)
   }
 
-  get(key?: T, keepPair?: boolean) {
+  /**
+   * If `keepPair` is `true`, returns the Pair matching `key`.
+   * Otherwise, returns the value of that Pair's key.
+   */
+  get(key: unknown, keepPair?: boolean): any {
     const pair = findPair(this.items, key)
     return !keepPair && isPair(pair)
       ? isScalar(pair.key)
@@ -46,7 +50,7 @@ export class YAMLSet<T = unknown> extends YAMLMap<T, Scalar<null> | null> {
 
   set(key: T, value: boolean): void
 
-  /** Will throw; `value` must be boolean */
+  /** @deprecated Will throw; `value` must be boolean */
   set(key: T, value: null): void
   set(key: T, value: boolean | null) {
     if (typeof value !== 'boolean')

src/visit.ts

@@ -320,7 +320,7 @@ function replaceNode(
     if (key === 'key') parent.key = node
     else parent.value = node
   } else if (isDocument(parent)) {
-    parent.contents = node
+    parent.contents = node as Node
   } else {
     const pt = isAlias(parent) ? 'alias' : 'scalar'
     throw new Error(`Cannot replace node with ${pt} parent`)