FocusDemo.tsx

const FocusDemo = () => {
    // with callback
    const [inputRef, setInputFocus] = useFocus<HTMLInputElement>();

   // on component mount
   const ref = useAutoFocus<HTMLInputElement>();

    return (
        <> 
            <button onClick={setInputFocus} >
               Focus
            </button>
            <input ref={inputRef} />
            <input ref={autoFocusRef} />
        </>
    )
    
}

const useAutoFocus = <T extends HTMLElement>() => {
  return useCallback((node: T) => {
    if (node) {
      node.focus();
    }
  }, []);
};

const useFocus = <T extends HTMLElement>() => {
    const htmlElRef = useRef<T>(null)
    const setFocus = () => {htmlElRef.current &&  htmlElRef.current.focus()}

    return [ htmlElRef, setFocus ];
}