From 65005df5737d25dfb2262bca4e0f1dfd1b7a3907 Mon Sep 17 00:00:00 2001
From: Mohd Ahmad <61076339+MohdAhmad1@users.noreply.github.com>
Date: Wed, 17 Jul 2024 19:30:29 +0530
Subject: [PATCH] update row factory api to include children and expanded
 element

---
 .../row-dragging/RowDraggingExample.tsx       |  15 +-
 package/DataTable.tsx                         |  15 +-
 package/DataTableRow.tsx                      | 180 +++++++++++++-----
 package/types/DataTableProps.ts               |  17 +-
 4 files changed, 157 insertions(+), 70 deletions(-)

diff --git a/app/examples/row-dragging/RowDraggingExample.tsx b/app/examples/row-dragging/RowDraggingExample.tsx
index 65ae9ebe..f2c5b8b1 100644
--- a/app/examples/row-dragging/RowDraggingExample.tsx
+++ b/app/examples/row-dragging/RowDraggingExample.tsx
@@ -53,15 +53,16 @@ export function RowDraggingExample() {
               height={400}
               withTableBorder
               withColumnBorders
-              rowFactory={({ record, index }) => (
+              rowFactory={({ record, index, rowProps, children }) => (
                 <Draggable key={record.id} draggableId={record.id} index={index}>
                   {(provided) => (
-                    <Table.Tr ref={provided.innerRef} {...provided.draggableProps} {...provided.dragHandleProps}>
-                      {columns.map((column) => (
-                        <Table.Td key={column.accessor}>
-                          {column.render ? column.render(record) : record[column.accessor as keyof RecordData]}
-                        </Table.Td>
-                      ))}
+                    <Table.Tr
+                      ref={provided.innerRef}
+                      {...rowProps}
+                      {...provided.draggableProps}
+                      {...provided.dragHandleProps}
+                    >
+                      {children}
                     </Table.Tr>
                   )}
                 </Draggable>
diff --git a/package/DataTable.tsx b/package/DataTable.tsx
index eb053f44..83d556fa 100644
--- a/package/DataTable.tsx
+++ b/package/DataTable.tsx
@@ -368,6 +368,7 @@ export function DataTable<T>({
                   const isSelected = selectedRecordIds?.includes(recordId) || false;
 
                   let handleSelectionChange: React.MouseEventHandler | undefined;
+
                   if (onSelectedRecordsChange && selectedRecords) {
                     handleSelectionChange = (e) => {
                       if (e.nativeEvent.shiftKey && lastSelectionChangeIndex !== null) {
@@ -397,10 +398,8 @@ export function DataTable<T>({
                       setLastSelectionChangeIndex(index);
                     };
                   }
-                  const RowWrapperComponent = ({ children }: { children: React.ReactNode }) =>
-                    rowWrapper ? rowWrapper(children, record, index) : children;
 
-                  const rowContent = (
+                  return (
                     <DataTableRow<T>
                       key={recordId as React.Key}
                       record={record}
@@ -431,21 +430,15 @@ export function DataTable<T>({
                       selectionColumnClassName={selectionColumnClassName}
                       selectionColumnStyle={selectionColumnStyle}
                       idAccessor={idAccessor as string}
+                      rowFactory={rowFactory}
                     />
                   );
-
-                  return rowFactory ? (
-                    <RowWrapperComponent key={recordId as React.Key}>
-                      {rowFactory({ record, index, children: rowContent })}
-                    </RowWrapperComponent>
-                  ) : (
-                    <RowWrapperComponent key={recordId as React.Key}>{rowContent}</RowWrapperComponent>
-                  );
                 })
               ) : (
                 <DataTableEmptyRow />
               )}
             </tbody>
+
             {effectiveColumns.some(({ footer }) => footer) && (
               <DataTableFooter<T>
                 ref={footerRef}
diff --git a/package/DataTableRow.tsx b/package/DataTableRow.tsx
index 7f758457..899c6d19 100644
--- a/package/DataTableRow.tsx
+++ b/package/DataTableRow.tsx
@@ -1,4 +1,4 @@
-import { TableTr, type CheckboxProps, type MantineColor, type MantineStyleProp } from '@mantine/core';
+import { MantineTheme, TableTr, type CheckboxProps, type MantineColor, type MantineStyleProp } from '@mantine/core';
 import clsx from 'clsx';
 import { DataTableRowCell } from './DataTableRowCell';
 import { DataTableRowExpansion } from './DataTableRowExpansion';
@@ -9,10 +9,12 @@ import type {
   DataTableCellClickHandler,
   DataTableColumn,
   DataTableDefaultColumnProps,
+  DataTableProps,
   DataTableRowClickHandler,
   DataTableSelectionTrigger,
 } from './types';
 import { CONTEXT_MENU_CURSOR, POINTER_CURSOR } from './utilityClasses';
+import { Fragment } from 'react';
 
 type DataTableRowProps<T> = {
   record: T;
@@ -49,7 +51,7 @@ type DataTableRowProps<T> = {
   selectionColumnClassName: string | undefined;
   selectionColumnStyle: MantineStyleProp | undefined;
   idAccessor: string;
-};
+} & Pick<DataTableProps<T>, 'rowFactory'>;
 
 export function DataTableRow<T>({
   record,
@@ -79,47 +81,11 @@ export function DataTableRow<T>({
   selectorCellShadowVisible,
   selectionColumnClassName,
   selectionColumnStyle,
+  rowFactory,
 }: Readonly<DataTableRowProps<T>>) {
-  return (
-    <>
-      <TableTr
-        className={clsx(
-          'mantine-datatable-row',
-          {
-            [POINTER_CURSOR]:
-              onClick || onDoubleClick || (expansion?.isExpandable({ record, index }) && expansion?.expandOnClick),
-          },
-          { [CONTEXT_MENU_CURSOR]: onContextMenu },
-          typeof className === 'function' ? className(record, index) : className
-        )}
-        data-selected={selectionChecked || undefined}
-        onClick={(e) => {
-          if (expansion) {
-            const { isExpandable, isRowExpanded, expandOnClick, expandRow, collapseRow } = expansion;
-            if (isExpandable({ record, index }) && expandOnClick) {
-              if (isRowExpanded(record)) {
-                collapseRow(record);
-              } else {
-                expandRow(record);
-              }
-            }
-          }
-          onClick?.({ event: e, record, index });
-        }}
-        onDoubleClick={onDoubleClick ? (e) => onDoubleClick({ event: e, record, index }) : undefined}
-        onContextMenu={onContextMenu ? (e) => onContextMenu({ event: e, record, index }) : undefined}
-        style={[
-          color || backgroundColor
-            ? (theme) => {
-                const colorValue = color?.(record, index);
-                const backgroundColorValue = backgroundColor?.(record, index);
-                return getRowCssVariables({ theme, color: colorValue, backgroundColor: backgroundColorValue });
-              }
-            : undefined,
-          style?.(record, index),
-        ]}
-        {...customAttributes?.(record, index)}
-      >
+  function TableCols() {
+    return (
+      <Fragment>
         {selectionVisible && (
           <DataTableRowSelectorCell<T>
             className={selectionColumnClassName}
@@ -135,6 +101,7 @@ export function DataTableRow<T>({
             getCheckboxProps={getSelectionCheckboxProps}
           />
         )}
+
         {columns.map(({ hidden, ...columnProps }, columnIndex) => {
           if (hidden) return null;
 
@@ -185,16 +152,129 @@ export function DataTableRow<T>({
             />
           );
         })}
+      </Fragment>
+    );
+  }
+
+  const expandedElement = expansion && (
+    <DataTableRowExpansion
+      colSpan={columns.filter(({ hidden }) => !hidden).length + (selectionVisible ? 1 : 0)}
+      open={expansion.isRowExpanded(record)}
+      content={expansion.content({ record, index })}
+      collapseProps={expansion.collapseProps}
+    />
+  );
+
+  const rowProps = getRowProps({
+    record,
+    index,
+    selectionChecked,
+    onClick,
+    onDoubleClick,
+    onContextMenu,
+    expansion,
+    customAttributes,
+    color,
+    backgroundColor,
+    className,
+    style,
+  });
+
+  if (rowFactory) {
+    return rowFactory({
+      record,
+      index,
+      rowProps,
+      children: <TableCols />,
+      expandedElement,
+    });
+  }
+
+  return (
+    <>
+      <TableTr {...rowProps}>
+        <TableCols />
       </TableTr>
 
-      {expansion && (
-        <DataTableRowExpansion
-          colSpan={columns.filter(({ hidden }) => !hidden).length + (selectionVisible ? 1 : 0)}
-          open={expansion.isRowExpanded(record)}
-          content={expansion.content({ record, index })}
-          collapseProps={expansion.collapseProps}
-        />
-      )}
+      {expandedElement}
     </>
   );
 }
+
+type GetRowPropsArgs<T> = Readonly<
+  Pick<
+    DataTableRowProps<T>,
+    | 'record'
+    | 'index'
+    | 'selectionChecked'
+    | 'onClick'
+    | 'onDoubleClick'
+    | 'onContextMenu'
+    | 'expansion'
+    | 'customAttributes'
+    | 'color'
+    | 'backgroundColor'
+    | 'className'
+    | 'style'
+  >
+>;
+
+export function getRowProps<T>({
+  record,
+  index,
+  selectionChecked,
+  onClick,
+  onDoubleClick,
+  onContextMenu,
+  expansion,
+  customAttributes,
+  color,
+  backgroundColor,
+  className,
+  style,
+}: GetRowPropsArgs<T>) {
+  return {
+    className: clsx(
+      'mantine-datatable-row',
+      {
+        [POINTER_CURSOR]:
+          onClick || onDoubleClick || (expansion?.isExpandable({ record, index }) && expansion?.expandOnClick),
+      },
+      { [CONTEXT_MENU_CURSOR]: onContextMenu },
+      typeof className === 'function' ? className(record, index) : className
+    ),
+
+    ['data-selected']: selectionChecked || undefined,
+
+    onClick: (e: React.MouseEvent<HTMLTableRowElement, MouseEvent>) => {
+      if (expansion) {
+        const { isExpandable, isRowExpanded, expandOnClick, expandRow, collapseRow } = expansion;
+        if (isExpandable({ record, index }) && expandOnClick) {
+          if (isRowExpanded(record)) {
+            collapseRow(record);
+          } else {
+            expandRow(record);
+          }
+        }
+      }
+      onClick?.({ event: e, record, index });
+    },
+    onDoubleClick: onDoubleClick
+      ? (e: React.MouseEvent<HTMLTableRowElement, MouseEvent>) => onDoubleClick({ event: e, record, index })
+      : undefined,
+    onContextMenu: onContextMenu
+      ? (e: React.MouseEvent<HTMLTableRowElement, MouseEvent>) => onContextMenu({ event: e, record, index })
+      : undefined,
+    style: [
+      color || backgroundColor
+        ? (theme: MantineTheme) => {
+            const colorValue = color?.(record, index);
+            const backgroundColorValue = backgroundColor?.(record, index);
+            return getRowCssVariables({ theme, color: colorValue, backgroundColor: backgroundColorValue });
+          }
+        : undefined,
+      style?.(record, index),
+    ],
+    ...(customAttributes?.(record, index) ?? {}),
+  };
+}
diff --git a/package/types/DataTableProps.ts b/package/types/DataTableProps.ts
index 4b9ce9c4..bf7ff588 100644
--- a/package/types/DataTableProps.ts
+++ b/package/types/DataTableProps.ts
@@ -1,4 +1,11 @@
-import type { MantineShadow, MantineStyleProp, ScrollAreaProps, StylesRecord, TableProps } from '@mantine/core';
+import type {
+  MantineShadow,
+  MantineStyleProp,
+  ScrollAreaProps,
+  StylesRecord,
+  TableProps,
+  TableTrProps,
+} from '@mantine/core';
 import type { DataTableCellClickHandler } from './DataTableCellClickHandler';
 import { DataTableColorProps } from './DataTableColorProps';
 import type { DataTableColumnProps } from './DataTableColumnProps';
@@ -200,7 +207,13 @@ export type DataTableProps<T = Record<string, unknown>> = {
    * a function that receives the current record and its index as arguments
    * and returns a React node representing the row.
    */
-  rowFactory?: (props: { record: T; index: number; children: React.ReactNode }) => React.ReactNode;
+  rowFactory?: (props: {
+    record: T;
+    index: number;
+    children: React.ReactNode;
+    rowProps: TableTrProps;
+    expandedElement?: React.ReactNode;
+  }) => React.ReactNode;
 
   /**
    * Optional function returning an object of custom attributes to be applied to each row in the table.