fix: react-aria crashes when using withFragment

src/components/table/Table.ts

@@ -9,8 +9,8 @@ import {
 import { withFragment } from '../../withFragment'
 
 export const Table = withFragment(NextTable, 'table');
-export const TableBody = NextTableBody;
-export const TableCell = NextTableCell;
-export const TableColumn = NextTableColumn;
-export const TableHeader = NextTableHeader;
-export const TableRow = NextTableRow;
+export const TableBody = withFragment(NextTableBody, 'tableBody');
+export const TableCell = withFragment(NextTableCell, 'tableCell');
+export const TableColumn = withFragment(NextTableColumn, 'tableColumn');
+export const TableHeader = withFragment(NextTableHeader, 'tableHeader');
+export const TableRow = withFragment(NextTableRow, 'tableRow');

src/context.ts

@@ -1,9 +1,15 @@
-import { TableProps } from "@nextui-org/react";
+import { TableBodyProps, TableCellProps, TableColumnProps, TableHeaderProps, TableProps } from "@nextui-org/react";
+import { TableRowProps } from "@nextui-org/table/dist/base/table-row";
 import { createContext } from "react";
 
 export interface FragmentUIContext {
   defaults: {
     table?: Partial<TableProps>;
+    tableBody?: Partial<TableBodyProps<unknown>>;
+    tableCell?: Partial<TableCellProps>;
+    tableColumn?: Partial<TableColumnProps<unknown>>;
+    tableHeader?: Partial<TableHeaderProps<unknown>>;
+    tableRow?: Partial<TableRowProps>;
   }
 }
 

src/withFragment.tsx

@@ -1,7 +1,16 @@
 import {  useContext } from "react";
 import { FragmentUIContext } from "./context";
 
-export const withFragment = <P extends object>(Component: React.ComponentType<P>, configId: keyof FragmentUIContext['defaults']) => (props: P) => {
-  const context = useContext(FragmentUIContext);
-  return <Component {...context['defaults'][configId] || {}} {...props} />
+export const withFragment = <P extends object>(Component: React.ComponentType<P>, configId: keyof FragmentUIContext['defaults']) => {
+  const ComponentWithContext = (props: P) => {
+    const context = useContext(FragmentUIContext);
+    return <Component {...context['defaults'][configId] || {}} {...props} />
+  };
+
+  // fix react-aria collection
+  if ('getCollectionNode' in Component) {
+    ComponentWithContext.getCollectionNode = Component.getCollectionNode;
+  }
+
+  return ComponentWithContext;
 };