fix(Table): move rowSelection logic to getDerivedStateFromProps

components/table/Table.tsx

@@ -33,6 +33,8 @@ import {
   PaginationConfig,
   PrepareParamsArgumentsReturn,
   ExpandIconProps,
+  WithStore,
+  CheckboxPropsCache,
 } from './interface';
 import Pagination from '../pagination';
 import Icon from '../icon';
@@ -104,7 +106,7 @@ const createComponents = (components: TableComponents = {}, prevComponents?: Tab
   };
 };
 
-export default class Table<T> extends React.Component<TableProps<T>, TableState<T>> {
+class Table<T> extends React.Component<TableProps<T>, TableState<T>> {
   static Column = Column;
 
   static ColumnGroup = ColumnGroup;
@@ -164,6 +166,23 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
       };
     }
 
+    if (nextProps.rowSelection && 'selectedRowKeys' in nextProps.rowSelection) {
+      nextProps.store.setState({
+        selectedRowKeys: nextProps.rowSelection.selectedRowKeys || [],
+      });
+    } else if (prevProps.rowSelection && !nextProps.rowSelection) {
+      nextProps.store.setState({
+        selectedRowKeys: [],
+      });
+    }
+    if ('dataSource' in nextProps && nextProps.dataSource !== prevProps.dataSource) {
+      nextProps.store.setState({
+        selectionDirty: false,
+      });
+    }
+    // https://github.com/ant-design/ant-design/issues/10133
+    nextProps.setCheckboxPropsCache({});
+
     if (nextProps.components !== prevProps.components) {
       const components = createComponents(nextProps.components, prevProps.components);
 
@@ -176,12 +195,6 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
     return nextState;
   }
 
-  CheckboxPropsCache: {
-    [key: string]: any;
-  };
-
-  store: Store;
-
   columns: ColumnProps<T>[];
 
   components: TableComponents;
@@ -222,38 +235,13 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
       prevProps: props,
       components: createComponents(props.components),
     };
-
-    this.CheckboxPropsCache = {};
-
-    this.store = createStore({
-      selectedRowKeys: getRowSelection(props).selectedRowKeys || [],
-      selectionDirty: false,
-    });
   }
 
-  componentDidUpdate(prevProps: Readonly<TableProps<T>>) {
+  componentDidUpdate() {
     const { props } = this;
 
     this.columns = props.columns || normalizeColumns(props.children as React.ReactChildren);
 
-    if (props.rowSelection && 'selectedRowKeys' in props.rowSelection) {
-      this.store.setState({
-        selectedRowKeys: props.rowSelection.selectedRowKeys || [],
-      });
-    } else if (prevProps.rowSelection && !props.rowSelection) {
-      this.store.setState({
-        selectedRowKeys: [],
-      });
-    }
-    if ('dataSource' in props && props.dataSource !== prevProps.dataSource) {
-      this.store.setState({
-        selectionDirty: false,
-      });
-    }
-
-    // https://github.com/ant-design/ant-design/issues/10133
-    this.CheckboxPropsCache = {};
-
     if (this.getSortOrderColumns(this.columns).length > 0) {
       const sortState = this.getSortStateFromColumns(this.columns);
       if (
@@ -284,16 +272,16 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
     }
     const key = this.getRecordKey(item, index);
     // Cache checkboxProps
-    if (!this.CheckboxPropsCache[key]) {
+    if (!this.props.checkboxPropsCache[key]) {
-      this.CheckboxPropsCache[key] = rowSelection.getCheckboxProps(item) || {};
+      this.props.checkboxPropsCache[key] = rowSelection.getCheckboxProps(item) || {};
-      const checkboxProps = this.CheckboxPropsCache[key];
+      const checkboxProps = this.props.checkboxPropsCache[key];
       warning(
         !('checked' in checkboxProps) && !('defaultChecked' in checkboxProps),
         'Table',
         'Do not set `checked` or `defaultChecked` in `getCheckboxProps`. Please use `selectedRowKeys` instead.',
       );
     }
-    return this.CheckboxPropsCache[key];
+    return this.props.checkboxPropsCache[key];
   };
 
   getDefaultSelection() {
@@ -502,7 +490,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
     return {
       ...custom,
       prefixCls,
-      store: this.store,
+      store: this.props.store,
       rowKey: this.getRecordKey(record, index),
     };
   };
@@ -511,7 +499,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
     const { selectWay, record, checked, changeRowKeys, nativeEvent } = selectionInfo;
     const rowSelection = getRowSelection(this.props);
     if (rowSelection && !('selectedRowKeys' in rowSelection)) {
-      this.store.setState({ selectedRowKeys });
+      this.props.store.setState({ selectedRowKeys });
     }
     const data = this.getFlatData();
     if (!rowSelection.onChange && !rowSelection[selectWay]) {
@@ -611,7 +599,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
 
     this.setState(newState, () => {
       this.scrollToFirstRow();
-      this.store.setState({
+      this.props.store.setState({
         selectionDirty: false,
       });
       const { onChange } = this.props;
@@ -632,8 +620,10 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
   handleSelect = (record: T, rowIndex: number, e: CheckboxChangeEvent) => {
     const checked = e.target.checked;
     const nativeEvent = e.nativeEvent;
-    const defaultSelection = this.store.getState().selectionDirty ? [] : this.getDefaultSelection();
+    const defaultSelection = this.props.store.getState().selectionDirty
-    let selectedRowKeys = this.store.getState().selectedRowKeys.concat(defaultSelection);
+      ? []
+      : this.getDefaultSelection();
+    let selectedRowKeys = this.props.store.getState().selectedRowKeys.concat(defaultSelection);
     const key = this.getRecordKey(record, rowIndex);
     const { pivot } = this.state;
     const rows = this.getFlatCurrentPageData();
@@ -667,7 +657,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
       }
 
       this.setState({ pivot: realIndex });
-      this.store.setState({
+      this.props.store.setState({
         selectionDirty: true,
       });
       this.setSelectedRowKeys(selectedRowKeys, {
@@ -685,7 +675,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
       }
 
       this.setState({ pivot: realIndex });
-      this.store.setState({
+      this.props.store.setState({
         selectionDirty: true,
       });
       this.setSelectedRowKeys(selectedRowKeys, {
@@ -703,7 +693,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
     const nativeEvent = e.nativeEvent;
     const key = this.getRecordKey(record, rowIndex);
     const selectedRowKeys = [key];
-    this.store.setState({
+    this.props.store.setState({
       selectionDirty: true,
     });
     this.setSelectedRowKeys(selectedRowKeys, {
@@ -717,8 +707,10 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
 
   handleSelectRow = (selectionKey: string, index: number, onSelectFunc: SelectionItemSelectFn) => {
     const data = this.getFlatCurrentPageData();
-    const defaultSelection = this.store.getState().selectionDirty ? [] : this.getDefaultSelection();
+    const defaultSelection = this.props.store.getState().selectionDirty
-    const selectedRowKeys = this.store.getState().selectedRowKeys.concat(defaultSelection);
+      ? []
+      : this.getDefaultSelection();
+    const selectedRowKeys = this.props.store.getState().selectedRowKeys.concat(defaultSelection);
     const changeableRowKeys = data
       .filter((item, i) => !this.getCheckboxPropsByItem(item, i).disabled)
       .map((item, i) => this.getRecordKey(item, i));
@@ -763,7 +755,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
         break;
     }
 
-    this.store.setState({
+    this.props.store.setState({
       selectionDirty: true,
     });
     // when select custom selection, callback selections[n].onSelect
@@ -804,7 +796,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
     }
     this.setState(newState, () => this.scrollToFirstRow());
 
-    this.store.setState({
+    this.props.store.setState({
       selectionDirty: false,
     });
 
@@ -1038,7 +1030,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
         <span onClick={stopPropagation}>
           <SelectionBox
             type={type}
-            store={this.store}
+            store={this.props.store}
             rowIndex={rowKey}
             onChange={handleChange}
             defaultSelection={this.getDefaultSelection()}
@@ -1087,7 +1079,7 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
         );
         selectionColumn.title = selectionColumn.title || (
           <SelectionCheckboxAll
-            store={this.store}
+            store={this.props.store}
             locale={locale}
             data={data}
             getCheckboxPropsByItem={this.getCheckboxPropsByItem}
@@ -1374,3 +1366,40 @@ export default class Table<T> extends React.Component<TableProps<T>, TableState<
     return <ConfigConsumer>{this.renderComponent}</ConfigConsumer>;
   }
 }
+
+function withStore(
+  WrappedComponent: typeof Table,
+): React.ComponentClass<Omit<TableProps<any>, keyof WithStore>> {
+  class Component<T> extends React.Component<TableProps<T>> {
+    store: Store;
+
+    CheckboxPropsCache: CheckboxPropsCache;
+
+    constructor(props: TableProps<T>) {
+      super(props);
+
+      this.CheckboxPropsCache = {};
+
+      this.store = createStore({
+        selectedRowKeys: getRowSelection(props).selectedRowKeys || [],
+        selectionDirty: false,
+      });
+    }
+
+    setCheckboxPropsCache = (cache: CheckboxPropsCache) => (this.CheckboxPropsCache = cache);
+
+    render() {
+      return (
+        <WrappedComponent<T>
+          {...this.props}
+          store={this.store}
+          checkboxPropsCache={this.CheckboxPropsCache}
+          setCheckboxPropsCache={this.setCheckboxPropsCache}
+        />
+      );
+    }
+  }
+  return Component;
+}
+
+export default withStore(Table);

components/table/__tests__/Table.test.js

@@ -61,7 +61,7 @@ describe('Table', () => {
     ];
     wrapper.setProps({ columns: newColumns });
 
-    expect(wrapper.instance().columns).toBe(newColumns);
+    expect(wrapper.dive().instance().columns).toBe(newColumns);
   });
 
   it('loading with Spin', async () => {

components/table/interface.tsx

@@ -148,7 +148,17 @@ export interface TableEventListeners {
   [name: string]: any; // https://github.com/ant-design/ant-design/issues/17245#issuecomment-504807714
 }
 
-export interface TableProps<T> {
+export interface CheckboxPropsCache {
+  [key: string]: any;
+}
+
+export interface WithStore {
+  store: Store;
+  checkboxPropsCache: CheckboxPropsCache;
+  setCheckboxPropsCache: (cache: CheckboxPropsCache) => void;
+}
+
+export interface TableProps<T> extends WithStore {
   prefixCls?: string;
   dropdownPrefixCls?: string;
   rowSelection?: TableRowSelection<T>;

components/table/util.tsx

@@ -63,7 +63,7 @@ export function normalizeColumns(elements: React.ReactChildren) {
     if (element.key) {
       column.key = element.key;
     }
-    if (element.type && (element.type as any).__ANT_TABLE_COLUMN_GROUP) {
+    if (element.props && (element.props as any).children) {
       column.children = normalizeColumns(column.children);
     }
     columns.push(column);