From a70c4bdcfae793f3c7cc3da388222880d4e4d622 Mon Sep 17 00:00:00 2001
From: Dmitry Artemov <d-artemov@yandex-team.ru>
Date: Fri, 25 Oct 2024 14:54:14 +0200
Subject: [PATCH] feat(BaseTable): add renderHeaderCellContent prop

---
 .../BaseHeaderCell/BaseHeaderCell.tsx         | 76 +++++++++++--------
 .../BaseHeaderRow/BaseHeaderRow.tsx           |  3 +
 src/components/BaseTable/BaseTable.tsx        |  4 +
 3 files changed, 53 insertions(+), 30 deletions(-)

diff --git a/src/components/BaseHeaderCell/BaseHeaderCell.tsx b/src/components/BaseHeaderCell/BaseHeaderCell.tsx
index d3bd4ca..21350be 100644
--- a/src/components/BaseHeaderCell/BaseHeaderCell.tsx
+++ b/src/components/BaseHeaderCell/BaseHeaderCell.tsx
@@ -22,6 +22,7 @@ export interface BaseHeaderCellProps<TData, TValue> {
         | ((header: Header<TData, TValue>, parentHeader?: Header<TData, unknown>) => string);
     header: Header<TData, TValue>;
     parentHeader?: Header<TData, unknown>;
+    renderHeaderCellContent?: (props: {header: Header<TData, TValue>}) => React.ReactNode;
     renderResizeHandle?: (props: BaseResizeHandleProps<TData, TValue>) => React.ReactNode;
     renderSortIndicator?: (props: BaseSortIndicatorProps<TData, TValue>) => React.ReactNode;
     resizeHandleClassName?: string;
@@ -38,6 +39,7 @@ export const BaseHeaderCell = <TData, TValue>({
     className: classNameProp,
     header,
     parentHeader,
+    renderHeaderCellContent,
     renderResizeHandle,
     renderSortIndicator,
     resizeHandleClassName,
@@ -54,6 +56,49 @@ export const BaseHeaderCell = <TData, TValue>({
 
     const rowSpan = header.isPlaceholder ? header.getLeafHeaders().length : 1;
 
+    const renderContent = () => {
+        if (renderHeaderCellContent) {
+            return renderHeaderCellContent({
+                header,
+            });
+        }
+
+        return (
+            <React.Fragment>
+                {header.column.getCanSort() ? (
+                    <BaseSort header={header}>
+                        {flexRender(header.column.columnDef.header, header.getContext())}{' '}
+                        {renderSortIndicator ? (
+                            renderSortIndicator({
+                                className: b('sort-indicator', sortIndicatorClassName),
+                                header,
+                            })
+                        ) : (
+                            <BaseSortIndicator
+                                className={b('sort-indicator', sortIndicatorClassName)}
+                                header={header}
+                            />
+                        )}
+                    </BaseSort>
+                ) : (
+                    flexRender(header.column.columnDef.header, header.getContext())
+                )}
+                {header.column.getCanResize() &&
+                    (renderResizeHandle ? (
+                        renderResizeHandle({
+                            className: b('resize-handle', resizeHandleClassName),
+                            header,
+                        })
+                    ) : (
+                        <BaseResizeHandle
+                            className={b('resize-handle', resizeHandleClassName)}
+                            header={header}
+                        />
+                    ))}
+            </React.Fragment>
+        );
+    };
+
     return (
         <th
             className={b('header-cell', getHeaderCellClassModes(header), className)}
@@ -64,36 +109,7 @@ export const BaseHeaderCell = <TData, TValue>({
             {...attributes}
             style={getCellStyles(header, attributes?.style)}
         >
-            {header.column.getCanSort() ? (
-                <BaseSort header={header}>
-                    {flexRender(header.column.columnDef.header, header.getContext())}{' '}
-                    {renderSortIndicator ? (
-                        renderSortIndicator({
-                            className: b('sort-indicator', sortIndicatorClassName),
-                            header,
-                        })
-                    ) : (
-                        <BaseSortIndicator
-                            className={b('sort-indicator', sortIndicatorClassName)}
-                            header={header}
-                        />
-                    )}
-                </BaseSort>
-            ) : (
-                flexRender(header.column.columnDef.header, header.getContext())
-            )}
-            {header.column.getCanResize() &&
-                (renderResizeHandle ? (
-                    renderResizeHandle({
-                        className: b('resize-handle', resizeHandleClassName),
-                        header,
-                    })
-                ) : (
-                    <BaseResizeHandle
-                        className={b('resize-handle', resizeHandleClassName)}
-                        header={header}
-                    />
-                ))}
+            {renderContent()}
         </th>
     );
 };
diff --git a/src/components/BaseHeaderRow/BaseHeaderRow.tsx b/src/components/BaseHeaderRow/BaseHeaderRow.tsx
index 7cb5c7b..97748e2 100644
--- a/src/components/BaseHeaderRow/BaseHeaderRow.tsx
+++ b/src/components/BaseHeaderRow/BaseHeaderRow.tsx
@@ -16,6 +16,7 @@ export interface BaseHeaderRowProps<TData, TValue = unknown>
         | ((headerGroup: HeaderGroup<TData>, parentHeaderGroup?: HeaderGroup<TData>) => string);
     headerGroup: HeaderGroup<TData>;
     parentHeaderGroup?: HeaderGroup<TData>;
+    renderHeaderCellContent?: BaseHeaderCellProps<TData, TValue>['renderHeaderCellContent'];
     renderResizeHandle?: (props: BaseResizeHandleProps<TData, TValue>) => React.ReactNode;
     renderSortIndicator: BaseHeaderCellProps<TData, TValue>['renderSortIndicator'];
     resizeHandleClassName?: BaseHeaderCellProps<TData, TValue>['resizeHandleClassName'];
@@ -34,6 +35,7 @@ export const BaseHeaderRow = <TData, TValue = unknown>({
     className: classNameProp,
     headerGroup,
     parentHeaderGroup,
+    renderHeaderCellContent,
     renderResizeHandle,
     renderSortIndicator,
     resizeHandleClassName,
@@ -65,6 +67,7 @@ export const BaseHeaderRow = <TData, TValue = unknown>({
                         className={cellClassName}
                         header={header as Header<TData, TValue>}
                         parentHeader={parentHeader}
+                        renderHeaderCellContent={renderHeaderCellContent}
                         renderResizeHandle={renderResizeHandle}
                         renderSortIndicator={renderSortIndicator}
                         resizeHandleClassName={resizeHandleClassName}
diff --git a/src/components/BaseTable/BaseTable.tsx b/src/components/BaseTable/BaseTable.tsx
index 88c7e4c..fdd8cd4 100644
--- a/src/components/BaseTable/BaseTable.tsx
+++ b/src/components/BaseTable/BaseTable.tsx
@@ -81,6 +81,8 @@ export interface BaseTableProps<TData, TScrollElement extends Element | Window =
     renderGroupHeader?: BaseRowProps<TData>['renderGroupHeader'];
     /** Function to override the default rendering of the entire group header row */
     renderGroupHeaderRowContent?: BaseRowProps<TData>['renderGroupHeaderRowContent'];
+    /** Function to override the default rendering of header cell content */
+    renderHeaderCellContent?: BaseHeaderRowProps<TData>['renderHeaderCellContent'];
     /** Function to override the default rendering of resize handles */
     renderResizeHandle?: BaseHeaderRowProps<TData>['renderResizeHandle'];
     /** Function to override the default rendering of sort indicators */
@@ -137,6 +139,7 @@ export const BaseTable = React.forwardRef(
             renderCustomRowContent,
             renderGroupHeader,
             renderGroupHeaderRowContent,
+            renderHeaderCellContent,
             renderResizeHandle,
             renderSortIndicator,
             resizeHandleClassName,
@@ -252,6 +255,7 @@ export const BaseTable = React.forwardRef(
                                 className={headerRowClassName}
                                 headerGroup={headerGroup}
                                 parentHeaderGroup={headerGroups[index - 1]}
+                                renderHeaderCellContent={renderHeaderCellContent}
                                 renderResizeHandle={renderResizeHandle}
                                 renderSortIndicator={renderSortIndicator}
                                 resizeHandleClassName={resizeHandleClassName}