From 0588b3a52fcee1575270b7a286f1b0d31f09acb9 Mon Sep 17 00:00:00 2001
From: kseniyakuzina <kseniyakuzina@yandex-team.ru>
Date: Tue, 13 Aug 2024 13:18:51 +0300
Subject: [PATCH] feat(Table): attributes customization

---
 src/components/Cell/Cell.tsx                 |  7 ++++
 src/components/DraggableRow/DraggableRow.tsx | 41 +++++++++++++-------
 src/components/HeaderCell/HeaderCell.tsx     | 13 +++++++
 src/components/HeaderRow/HeaderRow.tsx       | 17 +++++++-
 src/components/Row/Row.tsx                   | 30 ++++++++++----
 src/components/Table/Table.tsx               | 28 +++++++++++--
 6 files changed, 109 insertions(+), 27 deletions(-)

diff --git a/src/components/Cell/Cell.tsx b/src/components/Cell/Cell.tsx
index 66f7d2f..d56059b 100644
--- a/src/components/Cell/Cell.tsx
+++ b/src/components/Cell/Cell.tsx
@@ -10,6 +10,9 @@ export interface CellProps<TData>
     extends Omit<React.TdHTMLAttributes<HTMLTableCellElement>, 'className'> {
     cell?: CellType<TData, unknown>;
     className?: string | ((cell?: CellType<TData, unknown>) => string);
+    attributes?:
+        | React.TdHTMLAttributes<HTMLTableCellElement>
+        | ((cell?: CellType<TData, unknown>) => React.TdHTMLAttributes<HTMLTableCellElement>);
 }
 
 export const Cell = <TData,>({
@@ -17,17 +20,21 @@ export const Cell = <TData,>({
     children,
     className: classNameProp,
     style,
+    attributes: attributesProp,
     ...restProps
 }: CellProps<TData>) => {
     const className = React.useMemo(() => {
         return typeof classNameProp === 'function' ? classNameProp(cell) : classNameProp;
     }, [cell, classNameProp]);
 
+    const attributes = typeof attributesProp === 'function' ? attributesProp(cell) : attributesProp;
+
     return (
         <td
             className={b('cell', getCellClassModes(cell), className)}
             style={getCellStyles(cell, style)}
             {...restProps}
+            {...attributes}
         >
             {cell ? flexRender(cell.column.columnDef.cell, cell.getContext()) : children}
         </td>
diff --git a/src/components/DraggableRow/DraggableRow.tsx b/src/components/DraggableRow/DraggableRow.tsx
index c1872e0..c25f330 100644
--- a/src/components/DraggableRow/DraggableRow.tsx
+++ b/src/components/DraggableRow/DraggableRow.tsx
@@ -2,6 +2,7 @@ import React from 'react';
 
 import {useSortable} from '@dnd-kit/sortable';
 import {useForkRef} from '@gravity-ui/uikit';
+import type {Row as RowType} from '@tanstack/react-table';
 
 import {useDraggableRowDepth, useDraggableRowStyle} from '../../hooks';
 import type {RowProps} from '../Row';
@@ -14,7 +15,12 @@ export interface DraggableRowProps<TData, TScrollElement extends Element | Windo
 
 export const DraggableRow = React.forwardRef(
     <TData, TScrollElement extends Element | Window = HTMLDivElement>(
-        {getRowAttributes, row, style, ...restProps}: DraggableRowProps<TData, TScrollElement>,
+        {
+            attributes: attributesProp,
+            row,
+            style,
+            ...restProps
+        }: DraggableRowProps<TData, TScrollElement>,
         ref: React.Ref<HTMLTableRowElement>,
     ) => {
         const {setNodeRef, transform, transition, isDragging} = useSortable({
@@ -48,25 +54,30 @@ export const DraggableRow = React.forwardRef(
             nestingEnabled: enableNesting,
         });
 
-        const getDraggableRowDataAttributes = React.useCallback<
-            NonNullable<RowProps<TData>['getRowAttributes']>
-        >(
-            (draggableRow) => ({
-                ...getRowAttributes?.(draggableRow),
-                'data-key': draggableRow.id,
-                'data-depth': depth,
-                'data-draggable': true,
-                'data-dragging': isDragging,
-                'data-drag-active': isDragActive,
-                'data-expanded': isDragActive && isParent,
-            }),
-            [getRowAttributes, depth, isDragging, isDragActive, isParent],
+        const getDraggableRowDataAttributes = React.useCallback(
+            (draggableRow: RowType<TData>) => {
+                const attributes =
+                    typeof attributesProp === 'function'
+                        ? attributesProp(draggableRow)
+                        : attributesProp;
+
+                return {
+                    ...attributes,
+                    'data-key': draggableRow.id,
+                    'data-depth': depth,
+                    'data-draggable': true,
+                    'data-dragging': isDragging,
+                    'data-drag-active': isDragActive,
+                    'data-expanded': isDragActive && isParent,
+                };
+            },
+            [attributesProp, depth, isDragging, isDragActive, isParent],
         );
 
         return (
             <Row
                 ref={handleRowRef}
-                getRowAttributes={getDraggableRowDataAttributes}
+                attributes={getDraggableRowDataAttributes}
                 row={row}
                 style={draggableStyle}
                 {...restProps}
diff --git a/src/components/HeaderCell/HeaderCell.tsx b/src/components/HeaderCell/HeaderCell.tsx
index 7b8df44..e2ef382 100644
--- a/src/components/HeaderCell/HeaderCell.tsx
+++ b/src/components/HeaderCell/HeaderCell.tsx
@@ -20,6 +20,12 @@ export interface HeaderCellProps<TData, TValue> {
     renderSortIndicator?: (props: SortIndicatorProps<TData, TValue>) => React.ReactNode;
     resizeHandleClassName?: string;
     sortIndicatorClassName?: string;
+    attributes?:
+        | React.TdHTMLAttributes<HTMLTableCellElement>
+        | ((
+              header: Header<TData, TValue>,
+              parentHeader?: Header<TData, unknown>,
+          ) => React.TdHTMLAttributes<HTMLTableCellElement>);
 }
 
 export const HeaderCell = <TData, TValue>({
@@ -30,6 +36,7 @@ export const HeaderCell = <TData, TValue>({
     renderSortIndicator,
     resizeHandleClassName,
     sortIndicatorClassName,
+    attributes: attributesProp,
 }: HeaderCellProps<TData, TValue>) => {
     const className = React.useMemo(() => {
         return typeof classNameProp === 'function'
@@ -53,6 +60,11 @@ export const HeaderCell = <TData, TValue>({
 
     const rowSpan = header.isPlaceholder ? header.getLeafHeaders().length : 1;
 
+    const attributes =
+        typeof attributesProp === 'function'
+            ? attributesProp(header, parentHeader)
+            : attributesProp;
+
     return (
         <th
             className={b('header-cell', getHeaderCellClassModes(header), className)}
@@ -60,6 +72,7 @@ export const HeaderCell = <TData, TValue>({
             rowSpan={rowSpan > 1 ? rowSpan : undefined}
             onClick={header.column.getToggleSortingHandler()}
             style={getCellStyles(header)}
+            {...attributes}
         >
             {flexRender(header.column.columnDef.header, header.getContext())}{' '}
             {header.column.getCanSort() &&
diff --git a/src/components/HeaderRow/HeaderRow.tsx b/src/components/HeaderRow/HeaderRow.tsx
index 328c85f..c0d5c24 100644
--- a/src/components/HeaderRow/HeaderRow.tsx
+++ b/src/components/HeaderRow/HeaderRow.tsx
@@ -18,6 +18,13 @@ export interface HeaderRowProps<TData, TValue> {
     renderSortIndicator: HeaderCellProps<TData, TValue>['renderSortIndicator'];
     resizeHandleClassName?: HeaderCellProps<TData, TValue>['resizeHandleClassName'];
     sortIndicatorClassName: HeaderCellProps<TData, TValue>['sortIndicatorClassName'];
+    attributes?:
+        | React.HTMLAttributes<HTMLTableRowElement>
+        | ((
+              headerGroup: HeaderGroup<TData>,
+              parentHeaderGroup?: HeaderGroup<TData>,
+          ) => React.HTMLAttributes<HTMLTableRowElement>);
+    cellAttributes?: HeaderCellProps<TData, TValue>['attributes'];
 }
 
 export const HeaderRow = <TData, TValue>({
@@ -29,6 +36,8 @@ export const HeaderRow = <TData, TValue>({
     renderSortIndicator,
     resizeHandleClassName,
     sortIndicatorClassName,
+    attributes: attributesProp,
+    cellAttributes,
 }: HeaderRowProps<TData, TValue>) => {
     const className = React.useMemo(() => {
         return typeof classNameProp === 'function'
@@ -36,8 +45,13 @@ export const HeaderRow = <TData, TValue>({
             : classNameProp;
     }, [classNameProp, headerGroup, parentHeaderGroup]);
 
+    const attributes =
+        typeof attributesProp === 'function'
+            ? attributesProp(headerGroup, parentHeaderGroup)
+            : attributesProp;
+
     return (
-        <tr className={b('header-row', className)}>
+        <tr className={b('header-row', className)} {...attributes}>
             {headerGroup.headers.map((header) => (
                 <HeaderCell
                     key={header.column.id}
@@ -50,6 +64,7 @@ export const HeaderRow = <TData, TValue>({
                     renderSortIndicator={renderSortIndicator}
                     resizeHandleClassName={resizeHandleClassName}
                     sortIndicatorClassName={sortIndicatorClassName}
+                    attributes={cellAttributes}
                 />
             ))}
         </tr>
diff --git a/src/components/Row/Row.tsx b/src/components/Row/Row.tsx
index 2bf65d2..24a118d 100644
--- a/src/components/Row/Row.tsx
+++ b/src/components/Row/Row.tsx
@@ -16,9 +16,6 @@ export interface RowProps<TData, TScrollElement extends Element | Window = HTMLD
     getGroupTitle?: (row: RowType<TData>) => React.ReactNode;
     getIsCustomRow?: (row: RowType<TData>) => boolean;
     getIsGroupHeaderRow?: (row: RowType<TData>) => boolean;
-    getRowAttributes?: (
-        row: RowType<TData>,
-    ) => React.TdHTMLAttributes<HTMLTableRowElement> | undefined;
     groupHeaderClassName?: string;
     onClick?: (row: RowType<TData>, event: React.MouseEvent<HTMLTableRowElement>) => void;
     renderCustomRowContent?: (props: {
@@ -37,6 +34,10 @@ export interface RowProps<TData, TScrollElement extends Element | Window = HTMLD
     rowVirtualizer?: Virtualizer<TScrollElement, HTMLTableRowElement>;
     style?: React.CSSProperties;
     virtualItem?: VirtualItem<HTMLTableRowElement>;
+    attributes?:
+        | React.HTMLAttributes<HTMLTableRowElement>
+        | ((row: RowType<TData>) => React.HTMLAttributes<HTMLTableRowElement>);
+    cellAttributes?: CellProps<TData>['attributes'];
 }
 
 export const Row = React.forwardRef(
@@ -47,7 +48,6 @@ export const Row = React.forwardRef(
             getGroupTitle,
             getIsCustomRow,
             getIsGroupHeaderRow,
-            getRowAttributes,
             groupHeaderClassName,
             onClick,
             renderCustomRowContent,
@@ -57,6 +57,8 @@ export const Row = React.forwardRef(
             rowVirtualizer,
             style,
             virtualItem,
+            attributes: attributesProp,
+            cellAttributes,
         }: RowProps<TData, TScrollElement>,
         ref: React.Ref<HTMLTableRowElement>,
     ) => {
@@ -83,7 +85,11 @@ export const Row = React.forwardRef(
                         row,
                     })
                 ) : (
-                    <Cell className={cellClassName} colSpan={row.getVisibleCells().length}>
+                    <Cell
+                        className={cellClassName}
+                        colSpan={row.getVisibleCells().length}
+                        attributes={cellAttributes}
+                    >
                         {renderGroupHeader ? (
                             renderGroupHeader({
                                 className: b('group-header', groupHeaderClassName),
@@ -107,9 +113,19 @@ export const Row = React.forwardRef(
 
             return row
                 .getVisibleCells()
-                .map((cell) => <Cell key={cell.id} cell={cell} className={cellClassName} />);
+                .map((cell) => (
+                    <Cell
+                        key={cell.id}
+                        cell={cell}
+                        className={cellClassName}
+                        attributes={cellAttributes}
+                    />
+                ));
         };
 
+        const attributes =
+            typeof attributesProp === 'function' ? attributesProp(row) : attributesProp;
+
         return (
             <tr
                 ref={rowRef}
@@ -127,7 +143,7 @@ export const Row = React.forwardRef(
                 )}
                 onClick={handleClick}
                 data-index={virtualItem?.index}
-                {...getRowAttributes?.(row)}
+                {...attributes}
             >
                 {renderRowContent()}
             </tr>
diff --git a/src/components/Table/Table.tsx b/src/components/Table/Table.tsx
index 7e1f81d..c30d0e8 100644
--- a/src/components/Table/Table.tsx
+++ b/src/components/Table/Table.tsx
@@ -28,7 +28,6 @@ export interface TableProps<TData, TScrollElement extends Element | Window = HTM
     getGroupTitle?: RowProps<TData>['getGroupTitle'];
     getIsCustomRow?: RowProps<TData>['getIsCustomRow'];
     getIsGroupHeaderRow?: RowProps<TData>['getIsGroupHeaderRow'];
-    getRowAttributes?: RowProps<TData>['getRowAttributes'];
     groupHeaderClassName?: RowProps<TData>['groupHeaderClassName'];
     headerCellClassName?: HeaderRowProps<TData, unknown>['cellClassName'];
     headerClassName?: string;
@@ -48,6 +47,13 @@ export interface TableProps<TData, TScrollElement extends Element | Window = HTM
     table: TableType<TData>;
     withFooter?: boolean;
     withHeader?: boolean;
+    attributes?: React.TableHTMLAttributes<HTMLTableElement>;
+    headerAttributes?: React.HTMLAttributes<HTMLTableSectionElement>;
+    headerRowAttributes?: HeaderRowProps<TData, unknown>['attributes'];
+    headerCellAttributes?: HeaderRowProps<TData, unknown>['cellAttributes'];
+    bodyAttributes?: React.HTMLAttributes<HTMLTableSectionElement>;
+    rowAttributes?: RowProps<TData>['attributes'];
+    cellAttributes?: RowProps<TData>['cellAttributes'];
 }
 
 export const Table = React.forwardRef(
@@ -62,7 +68,6 @@ export const Table = React.forwardRef(
             footerRowClassName,
             getGroupTitle,
             getIsGroupHeaderRow,
-            getRowAttributes,
             headerCellClassName,
             headerClassName,
             headerRowClassName,
@@ -80,6 +85,13 @@ export const Table = React.forwardRef(
             table,
             withFooter,
             withHeader = true,
+            attributes,
+            headerAttributes,
+            headerRowAttributes,
+            headerCellAttributes,
+            bodyAttributes,
+            rowAttributes,
+            cellAttributes,
         }: TableProps<TData, TScrollElement>,
         ref: React.Ref<HTMLTableElement>,
     ) => {
@@ -104,9 +116,13 @@ export const Table = React.forwardRef(
                     ref={ref}
                     className={b({'with-row-virtualization': Boolean(rowVirtualizer)}, className)}
                     data-dragging-row-index={draggingRowIndex > -1 ? draggingRowIndex : undefined}
+                    {...attributes}
                 >
                     {headerGroups && (
-                        <thead className={b('header', {sticky: stickyHeader}, headerClassName)}>
+                        <thead
+                            className={b('header', {sticky: stickyHeader}, headerClassName)}
+                            {...headerAttributes}
+                        >
                             {headerGroups.map((headerGroup, index) => (
                                 <HeaderRow
                                     key={headerGroup.id}
@@ -118,6 +134,8 @@ export const Table = React.forwardRef(
                                     renderSortIndicator={renderSortIndicator}
                                     resizeHandleClassName={resizeHandleClassName}
                                     sortIndicatorClassName={sortIndicatorClassName}
+                                    attributes={headerRowAttributes}
+                                    cellAttributes={headerCellAttributes}
                                 />
                             ))}
                         </thead>
@@ -127,6 +145,7 @@ export const Table = React.forwardRef(
                         style={{
                             height: rowVirtualizer?.getTotalSize(),
                         }}
+                        {...bodyAttributes}
                     >
                         {(rowVirtualizer?.getVirtualItems() || rows).map((virtualItemOrRow) => {
                             const row = rowVirtualizer
@@ -138,7 +157,8 @@ export const Table = React.forwardRef(
                                 className: rowClassName,
                                 getGroupTitle,
                                 getIsGroupHeaderRow,
-                                getRowAttributes,
+                                attributes: rowAttributes,
+                                cellAttributes,
                                 onClick: onRowClick,
                                 renderGroupHeader,
                                 renderGroupHeaderRowContent,