feat: Add transformColumns for antd usage

Author: zombieJ

src/Table.tsx

@@ -101,6 +101,15 @@ export interface TableProps<RecordType extends DefaultRecordType>
   onRow?: GetComponentProps<RecordType>;
   onHeaderRow?: GetComponentProps<ColumnType<RecordType>[]>;
   emptyText?: React.ReactNode | (() => React.ReactNode);
+
+  // Internal
+  /**
+   * @private Internal usage, may remove by refactor. Should always use `columns` instead.
+   *
+   * !!! DO NOT USE IN PRODUCTION ENVIRONMENT !!!
+   */
+  // Used for antd table transform column with additional column
+  transformColumns: (columns: ColumnsType<RecordType>) => ColumnsType<RecordType>;
 }
 
 function Table<RecordType extends DefaultRecordType>(props: TableProps<RecordType>) {
@@ -127,6 +136,9 @@ function Table<RecordType extends DefaultRecordType>(props: TableProps<RecordTyp
     emptyText,
     onRow,
     onHeaderRow,
+
+    // Internal
+    transformColumns,
   } = props;
 
   const mergedData = data || EMPTY_DATA;
@@ -250,16 +262,19 @@ function Table<RecordType extends DefaultRecordType>(props: TableProps<RecordTyp
   // ====================== Column ======================
   const [componentWidth, setComponentWidth] = React.useState(0);
 
-  const [columns, flattenColumns] = useColumns({
-    ...props,
-    ...expandableConfig,
-    expandable: !!expandedRowRender,
-    expandedKeys: mergedExpandedKeys,
-    getRowKey,
-    onTriggerExpand,
-    expandIcon: mergedExpandIcon,
-    expandIconColumnIndex,
-  });
+  const [columns, flattenColumns] = useColumns(
+    {
+      ...props,
+      ...expandableConfig,
+      expandable: !!expandedRowRender,
+      expandedKeys: mergedExpandedKeys,
+      getRowKey,
+      onTriggerExpand,
+      expandIcon: mergedExpandIcon,
+      expandIconColumnIndex,
+    },
+    transformColumns,
+  );
 
   const columnContext = {
     columns,

src/hooks/useColumns.tsx

@@ -81,30 +81,33 @@ function warningFixed(flattenColumns: { fixed?: FixedType }[]) {
 /**
  * Parse `columns` & `children` into `columns`.
  */
-function useColumns<RecordType>({
-  prefixCls,
-  columns,
-  children,
-  expandable,
-  expandedKeys,
-  getRowKey,
-  onTriggerExpand,
-  expandIcon,
-  rowExpandable,
-  expandIconColumnIndex,
-}: {
-  prefixCls?: string;
-  columns?: ColumnsType<RecordType>;
-  children?: React.ReactNode;
-  expandable: boolean;
-  expandedKeys: Set<Key>;
-  getRowKey: GetRowKey<RecordType>;
-  onTriggerExpand: TriggerEventHandler<RecordType>;
-  expandIcon?: RenderExpandIcon<RecordType>;
-  rowExpandable?: (record: RecordType) => boolean;
-  expandIconColumnIndex?: number;
-}): [ColumnsType<RecordType>, ColumnType<RecordType>[]] {
-  const mergedColumns = React.useMemo<ColumnsType<RecordType>>(
+function useColumns<RecordType>(
+  {
+    prefixCls,
+    columns,
+    children,
+    expandable,
+    expandedKeys,
+    getRowKey,
+    onTriggerExpand,
+    expandIcon,
+    rowExpandable,
+    expandIconColumnIndex,
+  }: {
+    prefixCls?: string;
+    columns?: ColumnsType<RecordType>;
+    children?: React.ReactNode;
+    expandable: boolean;
+    expandedKeys: Set<Key>;
+    getRowKey: GetRowKey<RecordType>;
+    onTriggerExpand: TriggerEventHandler<RecordType>;
+    expandIcon?: RenderExpandIcon<RecordType>;
+    rowExpandable?: (record: RecordType) => boolean;
+    expandIconColumnIndex?: number;
+  },
+  transformColumns: (columns: ColumnsType<RecordType>) => ColumnsType<RecordType>,
+): [ColumnsType<RecordType>, ColumnType<RecordType>[]] {
+  const baseColumns = React.useMemo<ColumnsType<RecordType>>(
     () => columns || convertChildrenToColumns(children),
     [columns, children],
   );
@@ -113,7 +116,7 @@ function useColumns<RecordType>({
   const withExpandColumns = React.useMemo<ColumnsType<RecordType>>(() => {
     if (expandable) {
       const expandColIndex = expandIconColumnIndex || 0;
-      const prevColumn = mergedColumns[expandColIndex];
+      const prevColumn = baseColumns[expandColIndex];
 
       const expandColumn = {
         [INTERNAL_COL_DEFINE]: {
@@ -138,22 +141,27 @@ function useColumns<RecordType>({
       };
 
       // Insert expand column in the target position
-      const cloneColumns = mergedColumns.slice();
+      const cloneColumns = baseColumns.slice();
       cloneColumns.splice(expandColIndex, 0, expandColumn);
 
       return cloneColumns;
     }
-    return mergedColumns;
-  }, [expandable, mergedColumns, getRowKey, expandedKeys, expandIcon]);
+    return baseColumns;
+  }, [expandable, baseColumns, getRowKey, expandedKeys, expandIcon]);
 
-  const flattenColumns = React.useMemo(() => flatColumns(withExpandColumns), [withExpandColumns]);
+  const mergedColumns = React.useMemo(
+    () => (transformColumns ? transformColumns(withExpandColumns) : withExpandColumns),
+    [transformColumns, withExpandColumns],
+  );
+
+  const flattenColumns = React.useMemo(() => flatColumns(mergedColumns), [mergedColumns]);
 
   // Only check out of production since it's waste for each render
   if (process.env.NODE_ENV !== 'production') {
     warningFixed(flattenColumns);
   }
 
-  return [withExpandColumns, flattenColumns];
+  return [mergedColumns, flattenColumns];
 }
 
 export default useColumns;

tests/Table.spec.js

@@ -641,4 +641,16 @@ describe('Table.Basic', () => {
       ).render(),
     ).toMatchSnapshot();
   });
+
+  describe('internal api', () => {
+    it('transformColumns', () => {
+      const wrapper = mount(
+        createTable({
+          transformColumns: columns => [{ title: 'before' }, ...columns, { title: 'after' }],
+        }),
+      );
+
+      expect(wrapper.render()).toMatchSnapshot();
+    });
+  });
 });

tests/__snapshots__/Table.spec.js.snap

@@ -160,6 +160,77 @@ exports[`Table.Basic custom components renders fixed column and header correctly
 </div>
 `;
 
+exports[`Table.Basic internal api transformColumns 1`] = `
+<div
+  class="rc-table"
+>
+  <div
+    class="rc-table-content"
+  >
+    <table
+      style="table-layout: auto;"
+    >
+      <colgroup>
+        <col />
+        <col />
+        <col />
+      </colgroup>
+      <thead>
+        <tr>
+          <th
+            class="rc-table-cell"
+          >
+            before
+          </th>
+          <th
+            class="rc-table-cell"
+          >
+            Name
+          </th>
+          <th
+            class="rc-table-cell"
+          >
+            after
+          </th>
+        </tr>
+      </thead>
+      <tbody>
+        <tr
+          class="rc-table-row rc-table-row-level-0"
+        >
+          <td
+            class="rc-table-cell"
+          />
+          <td
+            class="rc-table-cell"
+          >
+            Lucy
+          </td>
+          <td
+            class="rc-table-cell"
+          />
+        </tr>
+        <tr
+          class="rc-table-row rc-table-row-level-0"
+        >
+          <td
+            class="rc-table-cell"
+          />
+          <td
+            class="rc-table-cell"
+          >
+            Jack
+          </td>
+          <td
+            class="rc-table-cell"
+          />
+        </tr>
+      </tbody>
+    </table>
+  </div>
+</div>
+`;
+
 exports[`Table.Basic renders colSpan correctly 1`] = `
 <div
   class="rc-table"