From 8ad1d99b601b31e3da7b0d56370c0a43c127d21a Mon Sep 17 00:00:00 2001
From: Henry Lin <henry.ys.lin@gmail.com>
Date: Wed, 31 May 2023 19:35:35 +0800
Subject: [PATCH] feat: add onStop option to watch

---
 src/__tests__/core.test.tsx | 23 +++++++++++++++++++++++
 src/core.ts                 | 16 ++++++++++++----
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/__tests__/core.test.tsx b/src/__tests__/core.test.tsx
index f2e8fec..0adc7c3 100644
--- a/src/__tests__/core.test.tsx
+++ b/src/__tests__/core.test.tsx
@@ -504,6 +504,29 @@ describe('watch', () => {
 
     expect(effectFn).toBeCalledTimes(3);
   });
+  it('calls onStop function', () => {
+    const counter = ref(1);
+    const effectFn = jest.fn();
+    const cleanupFn = jest.fn();
+    const stopFn = jest.fn();
+
+    const runner = watch(
+      counter,
+      () => {
+        effectFn(counter.value);
+        return cleanupFn;
+      },
+      { onStop: stopFn, immediate: true }
+    );
+
+    expect(stopFn).toBeCalledTimes(0);
+    expect(cleanupFn).toBeCalledTimes(0);
+
+    runner();
+
+    expect(stopFn).toBeCalledTimes(1);
+    expect(cleanupFn).toBeCalledTimes(1);
+  });
   it('cleans up properly', () => {
     const counter = ref(1);
     const effectFn = jest.fn();
diff --git a/src/core.ts b/src/core.ts
index 51aa597..89e4d0b 100644
--- a/src/core.ts
+++ b/src/core.ts
@@ -361,6 +361,7 @@ export type WatchCallback<V = any, OV = any> = (
 export interface WatchOptions<Immediate = boolean> extends DebuggerOptions {
   immediate?: Immediate;
   deep?: boolean;
+  onStop?: () => void;
 }
 
 export type WatchStopHandle = () => void;
@@ -432,6 +433,7 @@ export const watch: WatchOverloads = <
   const immediate = options?.immediate ?? false;
   const onTrack = options?.onTrack ?? undefined;
   const onTrigger = options?.onTrigger ?? undefined;
+  const onStop = options?.onStop ?? undefined;
 
   let getter: () => any;
   let isMultiSource = false;
@@ -498,6 +500,7 @@ export const watch: WatchOverloads = <
     allowRecurse: true,
     onTrack,
     onTrigger,
+    onStop,
   });
   if (immediate) {
     job();
@@ -514,12 +517,17 @@ export const watch: WatchOverloads = <
   return () => effect.effect.stop();
 };
 
+export interface UseWatchOptions<Immediate = boolean> extends DebuggerOptions {
+  immediate?: Immediate;
+  deep?: boolean;
+}
+
 interface UseWatchOverloads {
   // overload: array of multiple sources + cb
   <T extends MultiWatchSources, Immediate extends Readonly<boolean> = false>(
     sources: [...T],
     cb: WatchCallback<MapSources<T, false>, MapSources<T, Immediate>>,
-    options?: WatchOptions<Immediate>
+    options?: UseWatchOptions<Immediate>
   ): void;
   // overload: multiple sources w/ `as const`
   // watch([foo, bar] as const, () => {})
@@ -530,19 +538,19 @@ interface UseWatchOverloads {
   >(
     source: T,
     cb: WatchCallback<MapSources<T, false>, MapSources<T, Immediate>>,
-    options?: WatchOptions<Immediate>
+    options?: UseWatchOptions<Immediate>
   ): void;
   // overload: single source + cb
   <T, Immediate extends Readonly<boolean> = false>(
     source: WatchSource<T>,
     cb: WatchCallback<T, Immediate extends true ? T | undefined : T>,
-    options?: WatchOptions<Immediate>
+    options?: UseWatchOptions<Immediate>
   ): void;
   // overload: watching reactive object w/ cb
   <T extends object, Immediate extends Readonly<boolean> = false>(
     source: T,
     cb: WatchCallback<T, Immediate extends true ? T | undefined : T>,
-    options?: WatchOptions<Immediate>
+    options?: UseWatchOptions<Immediate>
   ): void;
 }