From d3563ba7b6a29de17820a08e80f714ed7306ca1a Mon Sep 17 00:00:00 2001
From: Alec Larson <1925840+aleclarson@users.noreply.github.com>
Date: Wed, 3 Jul 2024 15:37:29 -0400
Subject: [PATCH] feat(traverse): visitor can return a leave callback

---
 src/object/traverse.ts        | 24 ++++++++++++++----------
 tests/object/traverse.test.ts | 16 ++++++++++++++++
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/src/object/traverse.ts b/src/object/traverse.ts
index b8404905..06c04594 100644
--- a/src/object/traverse.ts
+++ b/src/object/traverse.ts
@@ -1,4 +1,4 @@
-import { isArray, isIterable, isPlainObject, last } from 'radashi'
+import { isArray, isFunction, isIterable, isPlainObject, last } from 'radashi'
 
 /**
  * Iterate an object's properties and those of any nested objects.
@@ -66,14 +66,14 @@ export function traverse(
     }
 
     context.path.push(key)
-    if (
-      visitor(
-        (context.value = value),
-        (context.key = key),
-        context.parent,
-        context,
-      ) === false
-    ) {
+    const result = visitor(
+      (context.value = value),
+      (context.key = key),
+      context.parent,
+      context,
+    )
+
+    if (result === false) {
       return (ok = false)
     }
 
@@ -87,6 +87,10 @@ export function traverse(
       !context.parents.includes(value)
     ) {
       traverse(value)
+
+      if (isFunction(result)) {
+        result()
+      }
     }
 
     context.path.pop()
@@ -147,7 +151,7 @@ export type TraverseVisitor = (
   key: keyof any,
   parent: object,
   context: TraverseContext,
-) => boolean | void
+) => (() => void) | boolean | void
 
 /**
  * The context object for the `traverse` function.
diff --git a/tests/object/traverse.test.ts b/tests/object/traverse.test.ts
index 6878a670..12f4d8ac 100644
--- a/tests/object/traverse.test.ts
+++ b/tests/object/traverse.test.ts
@@ -301,4 +301,20 @@ describe('traverse', () => {
     expect(paths).toEqual([['a'], ['a', 'b']])
     expect(parents).toEqual([[obj], [obj, obj.a]])
   })
+
+  test('visitor can return a leave callback', () => {
+    const obj = { a: { b: 2 } }
+    const visited: [keyof any, unknown][] = []
+    _.traverse(obj, (value, key) => {
+      visited.push([key, value])
+      return () => {
+        visited.push([key, value])
+      }
+    })
+    expect(visited).toEqual([
+      ['a', { b: 2 }],
+      ['b', 2],
+      ['a', { b: 2 }],
+    ])
+  })
 })