React Collapse component

Collapse.tsx

import type { TransitionEvent } from 'react'
import { useLayoutEffect, useRef, useState } from 'react'

const transition = 'height 200ms ease-out'

export type CollapseProps = {
  children: React.ReactNode

  /**
   * Controls whether or not children are visible.
   */
  open?: boolean

  /**
   * Determines if children should render `null` or not while closed.
   */
  lazy?: boolean

  /**
   * Instantly expand and collapse children.
   */
  instant?: boolean

  /**
   * Called when opening or closing.
   */
  onComplete?: Function
}

export function Collapse({
  children,
  instant,
  lazy,
  onComplete,
  open,
  ...restProps
}: CollapseProps) {
  const [renderChildren, setRenderChildren] = useState(lazy ? open : true)
  const ref = useRef<HTMLDivElement>(null)
  const firstRender = useRef(true)
  const instantRender = instant || firstRender.current

  function openCollapse() {
    const node = ref.current
    requestAnimationFrame(() => {
      node.style.height = node.scrollHeight + 'px'
    })
  }

  function closeCollapse() {
    const node = ref.current
    requestAnimationFrame(() => {
      node.style.height = node.offsetHeight + 'px'
      node.style.overflow = 'hidden'
      requestAnimationFrame(() => {
        node.style.height = '0px'
      })
    })
  }

  function handleComplete() {
    const node = ref.current
    node.style.overflow = open ? 'initial' : 'hidden'
    if (open) {
      requestAnimationFrame(() => {
        node.style.margin = null
        node.style.height = 'auto'
      })
    }
    if (!open && lazy) {
      node.style.margin = '0'
      setRenderChildren(false)
    }
    if (firstRender.current) {
      firstRender.current = false
    } else if (onComplete) {
      onComplete()
    }
  }

  function handleTransitionEnd(event: TransitionEvent<HTMLDivElement>) {
    const node = ref.current
    if (event.target === node && event.propertyName === 'height') {
      handleComplete()
    }
  }

  useLayoutEffect(() => {
    if (lazy) {
      if (open) {
        if (renderChildren) {
          openCollapse()
        } else {
          setRenderChildren(true)
        }
      } else {
        closeCollapse()
      }
    } else {
      if (open) {
        openCollapse()
      } else {
        closeCollapse()
      }
    }
  }, [open, lazy, renderChildren])

  useLayoutEffect(() => {
    if (instantRender) {
      handleComplete()
    }
  }, [open, instantRender])

  return (
    <div
      ref={ref}
      onTransitionEnd={handleTransitionEnd}
      style={{ transition: instantRender ? undefined : transition }}
      {...restProps}
    >
      {renderChildren ? children : null}
    </div>
  )
}

export default Collapse