feat: include keyPath in NodePath

pg-nano · Sep 9, 2024 · d877d9b · d877d9b
1 parent fabb903
commit d877d9b
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 11 deletions.
diff --git a/node.ts b/node.ts
@@ -15,7 +15,7 @@ export type NodeByTag<TNodeTag extends NodeTag> = Node extends infer TNode
   : never
 
 export class NodePath<TNodeTag extends NodeTag = NodeTag> {
-  constructor(readonly tag: TNodeTag, readonly node: NodeByTag<TNodeTag>, readonly parent: NodePath | null) {}
+  constructor(readonly tag: TNodeTag, readonly node: NodeByTag<TNodeTag>, readonly parent: NodePath | null, readonly keyPath: readonly (string | number)[]) {}
   isList(): this is NodePath<"List"> {
     return this.tag === "List"
   }

diff --git a/scripts/generateTypes.ts b/scripts/generateTypes.ts
@@ -474,7 +474,7 @@ export type NodeByTag<TNodeTag extends NodeTag> = Node extends infer TNode
   : never
 
 export class NodePath<TNodeTag extends NodeTag = NodeTag> {
-  constructor(readonly tag: TNodeTag, readonly node: NodeByTag<TNodeTag>, readonly parent: NodePath | null) {}
+  constructor(readonly tag: TNodeTag, readonly node: NodeByTag<TNodeTag>, readonly parent: NodePath | null, readonly keyPath: readonly (string | number)[]) {}
   ${Array.from(
     nodeTypes,
     (name) => `

diff --git a/walk.ts b/walk.ts
@@ -20,27 +20,32 @@ export function walk(
   root: any,
   callback: Walker | Visitor,
   parent: NodePath | null = null,
+  keyPath: readonly (string | number)[] = [],
 ) {
   if (typeof callback !== "function") {
     const visitor = callback
     callback = (path) => visitor[path.tag]?.(path as any)
   }
   if (Array.isArray(root)) {
-    root.forEach((node) => walk(node, callback))
+    root.forEach((node, index) => {
+      walk(node, callback, parent, [...keyPath, index])
+    })
   } else if (typeof root === "object" && root !== null) {
     const keys = Object.keys(root)
     if (keys.length === 1 && /^[A-Z]/.test(keys[0])) {
       const tag = keys[0] as NodeTag
-      const path = new NodePath(tag, root[tag], parent)
-      const result = callback(path)
-      if (result === false) {
+      const node = root[tag]
+      const path = new NodePath(tag, node, parent, keyPath)
+      if (callback(path) === false) {
         return
       }
-      parent = path
-    }
-    for (const key in root) {
-      const value = root[key]
-      walk(value, callback)
+      for (const key in node) {
+        walk(node[key], callback, path, [...keyPath, key])
+      }
+    } else {
+      for (const key of keys) {
+        walk(root[key], callback, parent, [...keyPath, key])
+      }
     }
   }
 }