🌭

viteの開発環境でzustandのstoreの実装を変えた時にstateを保持してhot module reloadしたい

に公開
Zustand
hmr
tech

やりたかったことはタイトルの通り。
少し複雑なstoreを作っていて、実装を変える度にstateがリセットされて状態を復元するのが大変だったので、stateを保持したまま実装だけ読み替えてほしかった。

近しいIssueは以下だけど用途に合わなかったので、参考にしつつ少し改修。

https://github.com/pmndrs/zustand/issues/934
https://github.com/pmndrs/zustand/discussions/827

結論

これを定義して

hmr.ts
import type { UseBoundStore } from 'zustand';

// eslint-disable-next-line @typescript-eslint/no-explicit-any
export const zustandHmrFix = (name: string, useStore: UseBoundStore<any>) => {
  if (!import.meta.hot || import.meta.env.VITEST) {
    return;
  }

  const savedState = import.meta.hot!.data[name];
  if (savedState) {
    const newState = { ...useStore.getState(), ...JSON.parse(savedState) };
    useStore.setState(newState, true);
  }

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  useStore.subscribe((state: any) => {
    import.meta.hot!.data[name] = JSON.stringify(state);
  });

  // 以下は多分不要(後述)
  // import.meta.hot!.accept((newModule) => {
  //   if (!newModule) {
  //     return;
  //   }
  //   const savedState = import.meta.hot!.data[name];
  //   if (savedState) {
  //     const newState = { ...useStore.getState(), ...JSON.parse(savedState) };
  //     useStore.setState(newState, true);
  //   }
  // });
};

store定義側でこんな感じで適用

import { create } from 'zustand';

const useXxxStore = create(...)

// NOTE: 第一引数をキーとしてstateを保持するため複数storeの場合は一意な名前にする必要あり
zustandHmrFix('xxxStore', useXxxStore);

型定義をanyじゃなくてgenericにしたかったけど色々な都合であきらめた。開発用なので問題無しとする。

説明

開発環境のみ有効化したいのでHMRが有効な場合のみ適用する。ただしvitestでwatchしてる場合もhotが定義されてるっぽいので別途除外。

  if (!import.meta.hot || import.meta.env.VITEST) {
    return;
  }

↓のsubscribeで保持されたdataが存在していたら新しいstoreにsetする。
(dataが無い場合は一度もsubscribeが発生していない=初期状態で良いので何もしない)

  const savedState = import.meta.hot!.data[name];
  if (savedState) {
    // useStore.getState()で新しい実装のaction関数を取得 + stateはsavedStateの値で上書き
    const newState = { ...useStore.getState(), ...JSON.parse(savedState) };

    // 再読み込みされたら関連コンポーネントも再描画されて欲しいので第二引数はtrue
    // 用途によってはfalseでも良いかも
    useStore.setState(newState, true);
  }

stateが変更されたらimport.meta.hot.dataに格納する。JSON.stringifyはvalueがシリアライズできない場合対応するkeyを出力しないため、これでstateのみを保持できる。

  useStore.subscribe((state: any) => {
    import.meta.hot!.data[name] = JSON.stringify(state);
  });


(chromeのconsoleで確認)

補足

Issueのコメントを見ると import.meta.hot!.accept(...) でもstateを更新していたけど、create()とzustandHmrFix()を同じファイルで実行しておけば、ファイルを保存した時点で自動的にzustandHmrFixも再実行されるため、↓の定義があると二重に実行されるだけなので不要と判断した。

例えば(あるかわからないけど)keyの一元管理とかのためにzustandHmrFixを別ファイルで実行してるような場合は、storeのファイルを保存してもzustandHmrFix自体は再実行されない可能性がありそうなので、そういう場合は書必要になるのかもしれない。

  import.meta.hot!.accept((newModule) => {
    if (!newModule) {
      return;
    }
    const savedState = import.meta.hot!.data[name];
    if (savedState) {
      const newState = { ...useStore.getState(), ...JSON.parse(savedState) };
      useStore.setState(newState, true);
    }
  });

おわり

viteの // @refresh reset が効いたら嬉しかったんだけど、対象がコンポーネントだけなのか今回の用途には使えなかった

https://leapcell.medium.com/beyond-hmr-understanding-reacts-fast-refresh-d6d80ef0fe4e

Discussion