如何为支持可选选择器的React Hook标注TypeScript返回值类型

我自己实现了一个基于Ref存储状态、采用发布订阅模式的简易状态管理器，当前useStore仅支持传入选择器函数返回状态切片，但无法正确标注类型——希望实现：不传选择器时返回完整状态对象，传入选择器时返回对应切片类型。

现有可运行代码：

function createStore<Store>(initialState: Store) {
  const useStoreData = () => {
    const store = useRef<Store>(initialState);
    const subscribers = useRef(new Set<() => void>());
    const get = () => store.current;
    const set = (value: Partial<Store>) => {
      store.current = { ...store.current, ...value };
      subscribers.current.forEach((cb) => cb());
    }

    const subscribe = (callback: () => void) => {
      subscribers.current.add(callback);
      return () => subscribers.current.delete(callback);
    }

    return { get, set, subscribe };
  }

  const StoreContext = createContext<ReturnType<typeof useStoreData> | null>(null);

  function Provider({ children }: { children: React.ReactNode }) {
    return (
      <StoreContext.Provider value={useStoreData()}>
        {children}
      </StoreContext.Provider>
    );
  }

  function useStore<SelectorOutput>(
    selector: (store: Store) => SelectorOutput
  ): [SelectorOutput, (value: Partial<Store>) => void] {
    const store = useContext(StoreContext);
    if (!store) {
      throw new Error("Store not found");
    }

    const [state, setState] = useState(selector(store.get()));
    useEffect(() => {
      return store.subscribe(() => setState(selector(store.get())));
    }, [store, selector]);

    return [state, store.set];
  }

  return {
    Provider,
    useStore
  };
}

// 现有用法类型推导正常
const { Provider, useStore } = createStore({ first: '', last: '' })
function DisplayValue() {
  const [first] = useStore((state) => state.first);
} 

但以下用法无法正确推导类型：

function DisplayAll() {
  const [state] = useStore() // 期望类型为 { first: string, last: string }
}

要实现目标类型推导，核心是让TypeScript识别useStore的两种调用场景，推荐两种实现方式：

方法1：使用函数重载 #

给useStore定义两个调用签名，分别对应“不传选择器”和“传入选择器”的场景：

// 重载1：不传选择器，返回完整Store类型
function useStore(): [Store, (value: Partial<Store>) => void];
// 重载2：传入选择器，返回选择器输出类型
function useStore<SelectorOutput>(selector: (store: Store) => SelectorOutput): [SelectorOutput, (value: Partial<Store>) => void];
// 通用实现逻辑
function useStore<SelectorOutput = Store>(selector?: (store: Store) => SelectorOutput): [SelectorOutput, (value: Partial<Store>) => void] {
  const store = useContext(StoreContext);
  if (!store) {
    throw new Error("Store not found");
  }

  // 处理默认逻辑：不传选择器时返回完整状态
  const usedSelector = selector ?? ((s) => s as unknown as SelectorOutput);

  const [state, setState] = useState(usedSelector(store.get()));
  useEffect(() => {
    return store.subscribe(() => setState(usedSelector(store.get())));
  }, [store, usedSelector]);

  return [state, store.set];
}

方法2：可选参数+泛型默认值 #

将选择器参数设为可选，同时给泛型SelectorOutput设置默认值为Store，并提供默认的选择器函数：

function useStore<SelectorOutput = Store>(
  selector: (store: Store) => SelectorOutput = (s) => s as SelectorOutput
): [SelectorOutput, (value: Partial<Store>) => void] {
  const store = useContext(StoreContext);
  if (!store) {
    throw new Error("Store not found");
  }

  const [state, setState] = useState(selector(store.get()));
  useEffect(() => {
    return store.subscribe(() => setState(selector(store.get())));
  }, [store, selector]);

  return [state, store.set];
}

额外优化 #

原代码中useStore返回的set函数类型标注有误，已修正为(value: Partial<Store>) => void，匹配实际的状态合并逻辑。

验证效果 #

修改后两种用法均能正确推导类型：

// 传入选择器，推导类型为string
function DisplayValue() {
  const [first] = useStore((state) => state.first);
} 

// 不传选择器，推导类型为{ first: string, last: string }
function DisplayAll() {
  const [state] = useStore();
}

内容的提问来源于stack exchange，提问作者joshuaaron