import React from 'react' ; import HeaderRow , { HeaderRowElement } from './HeaderRow' ; import { TableProps } from './Table' ; import { TableHeaderElement } from './TableHeader' ; import { isComponentType } from '@leafygreen-ui/lib' ; type TableHeaderProps < Shape > = Pick < TableProps < Shape > , 'columns' > ; function TableHead < Shape > ( { columns = [ ] } ) { const usingHeaderRow = React . useRef ( false ) ; function createCols ( array ) { return array . map ( ( child , index ) => { const tableHeaderCommonProps = { key : index , index , } ; if ( isComponentType < HeaderRowElement > ( child , 'HeaderRow' ) ) { usingHeaderRow . current = true ; const { children } = child ?. props ; return React . cloneElement ( child , { children : createCols ( React . Children . toArray ( children ) ) , } ) ; } if ( isComponentType < TableHeaderElement > ( child , 'TableHeader' ) ) { return React . cloneElement ( child , tableHeaderCommonProps ) ; } return child ; } ) ; } const columnArray = columns . type === React . Fragment || isComponentType < HeaderRowElement > ( columns , 'HeaderRow' ) ? React . Children . toArray ( ( columns as React . ReactElement ) . props . children ) : ( columns as Array < any > ) ; const cols = createCols ( columnArray ) ; if ( usingHeaderRow . current ) { return < thead > { cols } < /thead>; } return ( < thead > < HeaderRow > { cols } < /HeaderRow> < /thead>); } export default TableHead ;