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

eemeli · Apr 22, 2022 · 0691f44 · 0691f44
1 parent 3bed485
commit 0691f44
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 7 deletions.
diff --git a/src/doc/Document.ts b/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,

diff --git a/src/nodes/Node.ts b/src/nodes/Node.ts
@@ -13,6 +13,15 @@ export type Node<T = unknown> =
   | 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
   | Scalar.Parsed
@@ -32,7 +41,9 @@ 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 = <T = unknown>(node: any): node is Document<T> =>
+export const isDocument = <T extends Node = Node>(
+  node: any
+): node is Document<T> =>
   !!node && typeof node === 'object' && node[NODE_TYPE] === DOC
 
 export const isMap = <K = unknown, V = unknown>(

diff --git a/src/visit.ts b/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`)