Use transformation to provide isHeader property

src/elements/Table/Table.tsx

@@ -1,37 +1,19 @@
-import type { TableCellNode, TableNode } from '@prezly/story-content-format';
+import type { TableNode } from '@prezly/story-content-format';
 import classNames from 'classnames';
-import { Children, cloneElement, type HTMLAttributes, isValidElement } from 'react';
-
-import { isHeaderCell } from './TableContext';
+import type { HTMLAttributes } from 'react';
 
 interface Props extends HTMLAttributes<HTMLTableElement> {
     node: TableNode;
 }
 
 export function Table({ children, node }: Props) {
-    function isHeader(cell: TableCellNode): boolean {
-        return isHeaderCell(node, cell);
-    }
-
     return (
         <table
             className={classNames('prezly-slate-table', {
                 'prezly-slate-table--withBorders': node.border,
             })}
         >
-            <tbody>
-                {Children.map(children, (child) => {
-                    if (isValidElement(child)) {
-                        return cloneElement<{ isHeader?: (node: TableCellNode) => boolean }>(
-                            // @ts-expect-error I don't know how to type this :|
-                            child,
-                            { isHeader },
-                        );
-                    }
-
-                    return child;
-                })}
-            </tbody>
+            <tbody>{children}</tbody>
         </table>
     );
 }

src/elements/Table/TableCell.tsx

@@ -3,18 +3,16 @@ import classNames from 'classnames';
 import type { HTMLAttributes } from 'react';
 
 interface Props extends HTMLAttributes<HTMLTableCellElement> {
-    isHeader?: (node: TableCellNode) => boolean;
-    node: TableCellNode;
+    node: TableCellNode & { isHeader?: boolean };
 }
 
-export function TableCell({ children, isHeader, node }: Props) {
-    const isHeaderCell = isHeader?.(node);
-    const Cell = isHeaderCell ? 'th' : 'td';
+export function TableCell({ children, node }: Props) {
+    const Cell = node.isHeader ? 'th' : 'td';
 
     return (
         <Cell
             className={classNames('prezly-slate-table-cell', {
-                'prezly-slate-table-cell--header': isHeaderCell,
+                'prezly-slate-table-cell--header': node.isHeader,
             })}
             colSpan={node.colspan}
             rowSpan={node.rowspan}

src/elements/Table/TableRow.tsx

@@ -1,27 +1,10 @@
-import type { TableCellNode, TableRowNode } from '@prezly/story-content-format';
+import type { TableRowNode } from '@prezly/story-content-format';
 import type { HTMLAttributes } from 'react';
-import { Children, cloneElement, isValidElement } from 'react';
-
-import { TableCell as TableCellComponent } from './TableCell';
 
 interface Props extends HTMLAttributes<HTMLTableCellElement> {
-    isHeader?: (node: TableCellNode) => boolean;
     node: TableRowNode;
 }
 
-export function TableRow({ children, isHeader }: Props) {
-    return (
-        <tr className="prezly-slate-table-row">
-            {Children.map(children, (child) => {
-                if (isValidElement(child) && child.type === TableCellComponent) {
-                    // @ts-expect-error I don't know how to type this :|
-                    return cloneElement<{ isHeader?: (node: TableCellNode) => boolean }>(child, {
-                        isHeader,
-                    });
-                }
-
-                return child;
-            })}
-        </tr>
-    );
+export function TableRow({ children }: Props) {
+    return <tr className="prezly-slate-table-row">{children}</tr>;
 }

src/transformations/index.ts

@@ -1,2 +1,3 @@
 export { preventEmptyBlockCollapsing } from './preventEmptyBlockCollapsing';
 export { preventWhitespaceCollapsing } from './preventWhitepaceCollapsing';
+export { provideIsHeaderToTableCells } from './provideIsHeaderToTableCells';

src/transformations/provideIsHeaderToTableCells.ts

@@ -0,0 +1,30 @@
+import { type Node, TableNode } from '@prezly/story-content-format';
+
+export function provideIsHeaderToTableCells<T extends Node>(node: T): T {
+    if (TableNode.isTableNode(node)) {
+        const isHeaderFirstRow =
+            node.header &&
+            node.header.length === 1 &&
+            node.header[0] === TableNode.TableHeader.FIRST_ROW;
+
+        const isHeaderFirstColumn =
+            node.header &&
+            node.header.length === 1 &&
+            node.header[0] === TableNode.TableHeader.FIRST_COLUMN;
+
+        return {
+            ...node,
+            children: node.children.map((row, rowIndex) => ({
+                ...row,
+                children: row.children.map((cell, cellIndex) => ({
+                    ...cell,
+                    isHeader:
+                        (isHeaderFirstRow && rowIndex === 0) ||
+                        (isHeaderFirstColumn && cellIndex === 0),
+                })),
+            })),
+        };
+    }
+
+    return node;
+}