From 8bbfd347b162f57877d2c4e7ed9b217526f46451 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8F=B6=E6=9E=AB?=
 <7971419+crazyair@users.noreply.github.com>
Date: Thu, 30 Nov 2023 11:27:48 +0800
Subject: [PATCH] feat: useWatch support selector (#637)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: api

* feat: test

* feat: 代码优化

* feat: 添加显示声明类型支持

* feat: 下划线优化

* feat: 下划线优化

* feat: watchValue
---
 docs/demo/useWatch-selector.md      |  3 ++
 docs/examples/useWatch-selector.tsx | 37 ++++++++++++++++++++++
 src/useWatch.ts                     | 32 ++++++++++++++-----
 tests/useWatch.test.tsx             | 48 ++++++++++++++++++++++++++++-
 4 files changed, 112 insertions(+), 8 deletions(-)
 create mode 100644 docs/demo/useWatch-selector.md
 create mode 100644 docs/examples/useWatch-selector.tsx

diff --git a/docs/demo/useWatch-selector.md b/docs/demo/useWatch-selector.md
new file mode 100644
index 00000000..1c17ac8c
--- /dev/null
+++ b/docs/demo/useWatch-selector.md
@@ -0,0 +1,3 @@
+## useWatch-selector
+
+<code src="../examples/useWatch-selector.tsx"></code>
diff --git a/docs/examples/useWatch-selector.tsx b/docs/examples/useWatch-selector.tsx
new file mode 100644
index 00000000..7ecde689
--- /dev/null
+++ b/docs/examples/useWatch-selector.tsx
@@ -0,0 +1,37 @@
+import React from 'react';
+import Form, { Field } from 'rc-field-form';
+import Input from './components/Input';
+
+type FieldType = {
+  init?: string;
+  name?: string;
+  age?: number;
+};
+
+export default () => {
+  const [form] = Form.useForm<FieldType>();
+  const values = Form.useWatch(
+    values => ({ init: values.init, newName: values.name, newAge: values.age }),
+    { form, preserve: true },
+  );
+  console.log('values', values);
+  return (
+    <>
+      <Form form={form} initialValues={{ init: 'init', name: 'aaa' }}>
+        name
+        <Field name="name">
+          <Input />
+        </Field>
+        age
+        <Field name="age">
+          <Input />
+        </Field>
+        no-watch
+        <Field name="no-watch">
+          <Input />
+        </Field>
+        values:{JSON.stringify(values)}
+      </Form>
+    </>
+  );
+};
diff --git a/src/useWatch.ts b/src/useWatch.ts
index e5144083..127c04d5 100644
--- a/src/useWatch.ts
+++ b/src/useWatch.ts
@@ -76,6 +76,18 @@ function useWatch<TForm extends FormInstance>(
   form?: TForm | WatchOptions<TForm>,
 ): GetGeneric<TForm>;
 
+// ------- selector type -------
+function useWatch<TForm extends FormInstance, TSelected = unknown>(
+  selector: (values: GetGeneric<TForm>) => TSelected,
+  form?: TForm | WatchOptions<TForm>,
+): TSelected;
+
+function useWatch<ValueType = Store, TSelected = unknown>(
+  selector: (values: ValueType) => TSelected,
+  form?: FormInstance | WatchOptions<FormInstance>,
+): TSelected;
+// ------- selector type end -------
+
 function useWatch<TForm extends FormInstance>(
   dependencies: NamePath,
   form?: TForm | WatchOptions<TForm>,
@@ -86,8 +98,10 @@ function useWatch<ValueType = Store>(
   form?: FormInstance | WatchOptions<FormInstance>,
 ): ValueType;
 
-function useWatch(...args: [NamePath, FormInstance | WatchOptions<FormInstance>]) {
-  const [dependencies = [], _form = {}] = args;
+function useWatch(
+  ...args: [NamePath | ((values: Store) => any), FormInstance | WatchOptions<FormInstance>]
+) {
+  const [dependencies, _form = {}] = args;
   const options = isFormInstance(_form) ? { form: _form } : _form;
   const form = options.form;
 
@@ -125,8 +139,15 @@ function useWatch(...args: [NamePath, FormInstance | WatchOptions<FormInstance>]
       const { getFieldsValue, getInternalHooks } = formInstance;
       const { registerWatch } = getInternalHooks(HOOK_MARK);
 
+      const getWatchValue = (values: any, allValues: any) => {
+        const watchValue = options.preserve ? allValues : values;
+        return typeof dependencies === 'function'
+          ? dependencies(watchValue)
+          : getValue(watchValue, namePathRef.current);
+      };
+
       const cancelRegister = registerWatch((values, allValues) => {
-        const newValue = getValue(options.preserve ? allValues : values, namePathRef.current);
+        const newValue = getWatchValue(values, allValues);
         const nextValueStr = stringify(newValue);
 
         // Compare stringify in case it's nest object
@@ -137,10 +158,7 @@ function useWatch(...args: [NamePath, FormInstance | WatchOptions<FormInstance>]
       });
 
       // TODO: We can improve this perf in future
-      const initialValue = getValue(
-        options.preserve ? getFieldsValue(true) : getFieldsValue(),
-        namePathRef.current,
-      );
+      const initialValue = getWatchValue(getFieldsValue(), getFieldsValue(true));
 
       // React 18 has the bug that will queue update twice even the value is not changed
       // ref: https://github.com/facebook/react/issues/27213
diff --git a/tests/useWatch.test.tsx b/tests/useWatch.test.tsx
index 2c035975..843fe3a9 100644
--- a/tests/useWatch.test.tsx
+++ b/tests/useWatch.test.tsx
@@ -288,8 +288,29 @@ describe('useWatch', () => {
       const demo5 = Form.useWatch(['demo1', 'demo2', 'demo3', 'demo4', 'demo5'], form);
       const more = Form.useWatch(['age', 'name', 'gender'], form);
       const demo = Form.useWatch<string>(['demo']);
+
+      const values2 = Form.useWatch(values => ({ newName: values.name, newAge: values.age }), form);
+      const values3 = Form.useWatch<FieldType, { newName?: string }>(values => ({
+        newName: values.name,
+      }));
+
       return (
-        <>{JSON.stringify({ values, main, age, demo1, demo2, demo3, demo4, demo5, more, demo })}</>
+        <>
+          {JSON.stringify({
+            values,
+            main,
+            age,
+            demo1,
+            demo2,
+            demo3,
+            demo4,
+            demo5,
+            more,
+            demo,
+            values2,
+            values3,
+          })}
+        </>
       );
     };
 
@@ -457,4 +478,29 @@ describe('useWatch', () => {
 
     logSpy.mockRestore();
   });
+  it('selector', async () => {
+    const Demo: React.FC = () => {
+      const [form] = Form.useForm<{ name?: string }>();
+      const nameValue = Form.useWatch(values => values.name, form);
+      return (
+        <div>
+          <Form form={form}>
+            <Field name="name" initialValue="bamboo">
+              <Input />
+            </Field>
+          </Form>
+          <div className="values">{nameValue}</div>
+        </div>
+      );
+    };
+
+    const { container } = render(<Demo />);
+    await act(async () => {
+      await timeout();
+    });
+    expect(container.querySelector<HTMLDivElement>('.values')?.textContent).toEqual('bamboo');
+    const input = container.querySelectorAll<HTMLInputElement>('input');
+    await changeValue(input[0], 'bamboo2');
+    expect(container.querySelector<HTMLDivElement>('.values')?.textContent).toEqual('bamboo2');
+  });
 });