with-await-data.ts

import { createElement, lazy, useEffect, useMemo, ComponentType, ReactElement } from 'react';

export function withAwaitData<T extends {}, K extends keyof T, P extends T[K][]>(
  component: ComponentType<T>,
  fetcher: (...args: P) => Promise<void>,
  keys?: K[],
): (props: T) => ReactElement<T>;
export function withAwaitData<T extends {}, K extends keyof T, P extends T[K][]>(
  component: ComponentType<T>,
  fetcher: (...args: P) => Promise<void>,
  cleanup?: VoidFunction,
  keys?: K[],
): (props: T) => ReactElement<T>;
export function withAwaitData<T extends {}, K extends keyof T, P extends T[K][]>(
  component: ComponentType<T>,
  fetcher: (...args: P) => Promise<void>,
  cleanupOrKeys?: VoidFunction | K[],
  keysOrNothing?: K[],
): (props: T) => ReactElement<T> {
  let keys: K[];
  let cleanup: VoidFunction | undefined;

  if (typeof cleanupOrKeys === 'function') {
    cleanup = cleanupOrKeys;
    keys = keysOrNothing || [];
  } else {
    keys = cleanupOrKeys || [];
  }

  return (props: T) => {
    const args = useMemo(() => (keys.map((name) => props[name]) as P), [props]);
    const comp = useMemo(() => lazy(() => fetcher(...args).then(() => ({ default: component }))), args);

    useEffect(() => () => cleanup?.(), args);

    return createElement(comp as unknown as ComponentType<T>, props);
  };
}