ant-design · zombieJ · Jun 7, 2023 · Jun 1, 2023 · Jun 2, 2023 · Jun 2, 2023
diff --git a/components/transfer/__tests__/index.test.tsx b/components/transfer/__tests__/index.test.tsx
@@ -1,11 +1,11 @@
 import { fireEvent, render, waitFor } from '@testing-library/react';
-import React, { useState } from 'react';
+import type { DefaultRecordType } from 'rc-table/lib/interface';
+import React, { useCallback, useEffect, useState } from 'react';
 import type { SelectAllLabel, TransferProps } from '..';
 import Transfer from '..';
 import mountTest from '../../../tests/shared/mountTest';
 import rtlTest from '../../../tests/shared/rtlTest';
-
-const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
+import Button from '../../button';
 
 const listCommonProps: {
   dataSource: { key: string; title: string; disabled?: boolean }[];
@@ -264,6 +264,8 @@ describe('Transfer', () => {
   });
 
   it('should display the correct locale and ignore old API', () => {
+    const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
+
     const emptyProps = { dataSource: [], selectedKeys: [], targetKeys: [] };
     const locale = { notFoundContent: 'old1', searchPlaceholder: 'old2' };
     const newLocalProp = { notFoundContent: 'new1', searchPlaceholder: 'new2' };
@@ -621,4 +623,75 @@ describe('immutable data', () => {
     const { container } = render(<Transfer rowKey={(item) => item.id} dataSource={mockData} />);
     expect(container.firstChild).toMatchSnapshot();
   });
+
+  it('prevent error when reset data in some cases', () => {
+    const App = () => {
+      const [mockData, setMockData] = useState<DefaultRecordType[]>([]);
+      const [targetKeys, setTargetKeys] = useState<string[]>([]);
+
+      const getMock = () => {
+        const tempTargetKeys = [];
+        const tempMockData = [];
+        for (let i = 0; i < 2; i++) {
+          const data = {
+            key: i.toString(),
+            title: `content${i + 1}`,
+            description: `description of content${i + 1}`,
+            chosen: i % 2 === 0,
+          };
+          if (data.chosen) {
+            tempTargetKeys.push(data.key);
+          }
+          tempMockData.push(data);
+        }
+        setMockData(tempMockData);
+        setTargetKeys(tempTargetKeys);
+      };
+
+      useEffect(() => {
+        getMock();
+      }, []);
+
+      const handleChange = (newTargetKeys: string[]) => {
+        setTargetKeys(newTargetKeys);
+      };
+
+      const ButtonRender = useCallback(
+        () => <Button onClick={getMock}>Right button reload</Button>,
+        [getMock],
+      );
+
+      return (
+        <Transfer
+          dataSource={mockData}
+          operations={['to right', 'to left']}
+          targetKeys={targetKeys}
+          onChange={handleChange}
+          render={(item) => `test-${item}`}
+          footer={ButtonRender}
+        />
+      );
+    };
+
+    const { container } = render(<App />);
+    fireEvent.click(container.querySelector('.ant-transfer-list-header input[type="checkbox"]')!);
+    fireEvent.click(container.querySelector('.ant-transfer-operation .ant-btn')!);
+    expect(container.querySelectorAll('.ant-transfer-list')[1]).toBeTruthy();
+    expect(
+      container
+        .querySelectorAll('.ant-transfer-list')[1]
+        .querySelectorAll('.ant-transfer-list-content-item').length,
+    ).toBe(2);
+
+    fireEvent.click(
+      container.querySelectorAll('.ant-transfer-list-header input[type="checkbox"]')![1],
+    );
+    expect(container.querySelectorAll('.ant-transfer-list-header-selected')[1]).toContainHTML(
+      '2/2',
+    );
+    fireEvent.click(container.querySelector('.ant-transfer-list-footer .ant-btn')!);
+    expect(container.querySelectorAll('.ant-transfer-list-header-selected')[1]).toContainHTML(
+      '1/1',
+    );
+  });
 });
diff --git a/components/transfer/hooks/useData.ts b/components/transfer/hooks/useData.ts
@@ -0,0 +1,38 @@
+import * as React from 'react';
+import type { KeyWise, TransferProps } from '..';
+import { groupKeysMap } from '../../_util/transKeys';
+
+export default function useData<RecordType extends object>(
+  dataSource?: RecordType[],
+  rowKey?: TransferProps<RecordType>['rowKey'],
+  targetKeys: string[] = [],
+) {
+  const mergedDataSource = React.useMemo(
+    () =>
+      (dataSource || []).map((record: KeyWise<RecordType>) => {
+        if (rowKey) {
+          record = { ...record, key: rowKey(record) };
+        }
+        return record;
+      }),
+    [dataSource, rowKey],
+  );
+
+  const [leftDataSource, rightDataSource] = React.useMemo(() => {
+    const leftData: KeyWise<RecordType>[] = [];
+    const rightData: KeyWise<RecordType>[] = new Array(targetKeys.length);
+    const targetKeysMap = groupKeysMap(targetKeys);
+    mergedDataSource.forEach((record: KeyWise<RecordType>) => {
+      // rightData should be ordered by targetKeys
+      // leftData should be ordered by dataSource
+      if (targetKeysMap.has(record.key)) {
+        rightData[targetKeysMap.get(record.key)!] = record;
+      } else {
+        leftData.push(record);
+      }
+    });
+    return [leftData, rightData] as const;
+  }, [mergedDataSource, targetKeys, rowKey]);
+
+  return [mergedDataSource, leftDataSource, rightDataSource];
+}
diff --git a/components/transfer/hooks/useSelection.ts b/components/transfer/hooks/useSelection.ts
@@ -0,0 +1,57 @@
+import * as React from 'react';
+
+const EMPTY_KEYS: string[] = [];
+
+function filterKeys(keys: string[], dataKeys: Set<string>) {
+  const filteredKeys = keys.filter((key) => dataKeys.has(key));
+  return keys.length === filteredKeys.length ? keys : filteredKeys;
+}
+
+export default function useSelection<T extends { key: string }>(
+  leftDataSource: T[],
+  rightDataSource: T[],
+  selectedKeys: string[] = EMPTY_KEYS,
+): [
+  sourceSelectedKeys: string[],
+  targetSelectedKeys: string[],
+  setSourceSelectedKeys: React.Dispatch<React.SetStateAction<string[]>>,
+  setTargetSelectedKeys: React.Dispatch<React.SetStateAction<string[]>>,
+] {
+  // Prepare `dataSource` keys
+  const [leftKeys, rightKeys] = React.useMemo(
+    () => [
+      new Set(leftDataSource.map((src) => src.key)),
+      new Set(rightDataSource.map((src) => src.key)),
+    ],
+    [leftDataSource, rightDataSource],
+  );
+
+  // Selected Keys
+  const [sourceSelectedKeys, setSourceSelectedKeys] = React.useState(() =>
+    filterKeys(selectedKeys, leftKeys),
+  );
+  const [targetSelectedKeys, setTargetSelectedKeys] = React.useState(() =>
+    filterKeys(selectedKeys, rightKeys),
+  );
+
+  // Fill selected keys
+  React.useEffect(() => {
+    setSourceSelectedKeys(filterKeys(selectedKeys, leftKeys));
+    setTargetSelectedKeys(filterKeys(selectedKeys, rightKeys));
+  }, [selectedKeys]);
+
+  // Reset when data changed
+  React.useEffect(() => {
+    setSourceSelectedKeys(filterKeys(sourceSelectedKeys, leftKeys));
+    setTargetSelectedKeys(filterKeys(targetSelectedKeys, rightKeys));
+  }, [leftKeys, rightKeys]);
+
+  return [
+    // Keys
+    sourceSelectedKeys,
+    targetSelectedKeys,
+    // Updater
+    setSourceSelectedKeys,
+    setTargetSelectedKeys,
+  ];
+}
diff --git a/components/transfer/index.tsx b/components/transfer/index.tsx
@@ -1,6 +1,6 @@
 import classNames from 'classnames';
 import type { ChangeEvent, CSSProperties } from 'react';
-import React, { useCallback, useContext, useEffect, useMemo, useState } from 'react';
+import React, { useCallback, useContext } from 'react';
 import type { InputStatus } from '../_util/statusUtils';
 import { getMergedStatus, getStatusClassNames } from '../_util/statusUtils';
 import { groupDisabledKeysMap, groupKeysMap } from '../_util/transKeys';
@@ -12,13 +12,14 @@ import type { FormItemStatusContextProps } from '../form/context';
 import { FormItemInputContext } from '../form/context';
 import { useLocale } from '../locale';
 import defaultLocale from '../locale/en_US';
+import useData from './hooks/useData';
+import useSelection from './hooks/useSelection';
 import type { PaginationType } from './interface';
 import type { TransferListProps } from './list';
 import List from './list';
 import type { TransferListBodyProps } from './ListBody';
 import Operation from './operation';
 import Search from './search';
-
 import useStyle from './style';
 
 export type { TransferListProps } from './list';
@@ -108,9 +109,9 @@ const Transfer = <RecordType extends TransferItem = TransferItem>(
   props: TransferProps<RecordType>,
 ) => {
   const {
-    dataSource = [],
+    dataSource,
     targetKeys = [],
-    selectedKeys = [],
+    selectedKeys,
     selectAllLabels = [],
     operations = [],
     style = {},
@@ -147,20 +148,22 @@ const Transfer = <RecordType extends TransferItem = TransferItem>(
 
   const [wrapSSR, hashId] = useStyle(prefixCls);
 
-  const [sourceSelectedKeys, setSourceSelectedKeys] = useState<string[]>(() =>
-    selectedKeys.filter((key) => !targetKeys.includes(key)),
-  );
-
-  const [targetSelectedKeys, setTargetSelectedKeys] = useState<string[]>(() =>
-    selectedKeys.filter((key) => targetKeys.includes(key)),
+  // Fill record with `key`
+  const [mergedDataSource, leftDataSource, rightDataSource] = useData(
+    dataSource,
+    rowKey,
+    targetKeys,
   );
 
-  useEffect(() => {
-    if (props.selectedKeys) {
-      setSourceSelectedKeys(() => selectedKeys.filter((key) => !targetKeys.includes(key)));
-      setTargetSelectedKeys(() => selectedKeys.filter((key) => targetKeys.includes(key)));
-    }
-  }, [props.selectedKeys, props.targetKeys]);
+  // Get direction selected keys
+  const [
+    // Keys
+    sourceSelectedKeys,
+    targetSelectedKeys,
+    // Setters
+    setSourceSelectedKeys,
+    setTargetSelectedKeys,
+  ] = useSelection(leftDataSource, rightDataSource, selectedKeys);
 
   if (process.env.NODE_ENV !== 'production') {
     warning(
@@ -207,7 +210,7 @@ const Transfer = <RecordType extends TransferItem = TransferItem>(
 
   const moveTo = (direction: TransferDirection) => {
     const moveKeys = direction === 'right' ? sourceSelectedKeys : targetSelectedKeys;
-    const dataSourceDisabledKeysMap = groupDisabledKeysMap(dataSource);
+    const dataSourceDisabledKeysMap = groupDisabledKeysMap(mergedDataSource);
     // filter the disabled options
     const newMoveKeys = moveKeys.filter((key) => !dataSourceDisabledKeysMap.has(key));
     const newMoveKeysMap = groupKeysMap(newMoveKeys);
@@ -309,25 +312,6 @@ const Transfer = <RecordType extends TransferItem = TransferItem>(
     return listStyle || {};
   };
 
-  const [leftDataSource, rightDataSource] = useMemo(() => {
-    const leftData: KeyWise<RecordType>[] = [];
-    const rightData: KeyWise<RecordType>[] = new Array(targetKeys.length);
-    const targetKeysMap = groupKeysMap(targetKeys);
-    dataSource.forEach((record: KeyWise<RecordType>) => {
-      if (rowKey) {
-        record = { ...record, key: rowKey(record) };
-      }
-      // rightData should be ordered by targetKeys
-      // leftData should be ordered by dataSource
-      if (targetKeysMap.has(record.key)) {
-        rightData[targetKeysMap.get(record.key)!] = record;
-      } else {
-        leftData.push(record);
-      }
-    });
-    return [leftData, rightData] as const;
-  }, [dataSource, targetKeys, rowKey]);
-
   const formItemContext = useContext<FormItemStatusContextProps>(FormItemInputContext);
 
   const { hasFeedback, status } = formItemContext;