简单解读 zustand 源码（二）persist中间件源码解读

本次解读zustand的常用中间件persist源码，源码非常简单，主要学习设计思想首先找到esm/middleware.js

persist方法结构

可以看到如下代码：

// for compatibility of old code
const persistImpl = (config, baseOptions) => {
  if (
    "getStorage" in baseOptions ||
    "serialize" in baseOptions ||
    "deserialize" in baseOptions
  ) {
    if (process.env.NODE_ENV !== "production") {
      console.warn(
        "[DEPRECATED] `getStorage`, `serialize` and `deserialize` options are deprecated. Use `storage` option instead."
      );
    }
    return oldImpl(config, baseOptions);
  }
  return newImpl(config, baseOptions);
};
// rename
const persist = persistImpl;

这是兼容老版本的代码，我们只看新版本，即不支持getStorage、serialize和deserialize三个选项的新版本，newImpl

类型定义

首先可以看到代码的结构如下,接受两个参数，第一个参数是config，第二个参数为option，返回值是一个函数，这个函数类型和前文的create需要的类型是一样的。

const newImpl = (config, baseOptions) => (set, get, api) => {}

第一个参数config

查看zustand的类型定义，可以看到如下代码，可以看到config和返回值的类型是一样的，在用法上也可以体现，需要把本来传入给create的参数当作第一个参数传给persist

type PersistImpl = <T>(
  storeInitializer: StateCreator<T, [], []>,
  options: PersistOptions<T, T>
) => StateCreator<T, [], []>

第二个参数option

详情可以查看persist文档 (opens in a new tab)，在解读代码的时候也会接触到option

源码解读

初始化变量

首先初始化了几个变量

先覆盖掉默认的option
其次初始化hasHydrated为false，表示还没有把本地存在的state hydrate
定义hydrate时需要执行的listener -> hydrationListeners
定义hydrate结束的listener -> finishHydrationListeners

  let options = {
    // default options , storage === localStorage etc...
    storage: createJSONStorage(() => localStorage),
    // partialize can remove the properties that should not be stored , such as computed properties, functions, transient properties , private properties 
    // default does not remove any properties
    partialize: (state) => state,
    version: 0,
    merge: (persistedState, currentState) => ({
      ...currentState,
      //use persistedState first, use to load localStorage
      ...persistedState,
    }),
    ...baseOptions,
  };
  let hasHydrated = false;
  const hydrationListeners = /* @__PURE__ */ new Set();
  const finishHydrationListeners = /* @__PURE__ */ new Set();

判断storage是否存在

不存在就报错，没太多可说的

let storage = options.storage;
  // check if storage exists, config has same type as the return type of persist func
  if (!storage) {
    return config(
      (...args) => {
        //if not exists , do noting , and return 
        console.warn(
          `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`
        );
        set(...args);
      },
      get,
      api
    );
  }

劫持setState，加入同步写入storage逻辑

定义了新的setState代码，使其能够在setState的时候同时执行setItem写入本地。其中partialize方法用来过滤掉一些不需要放入本地的变量，如果不传那么不做任何改动

// use to set storage
  const setItem = () => {
    const state = options.partialize({ ...get() });
    return storage.setItem(options.name, {
      state,
      version: options.version,
    });
  };
  // api const api = { setState, getState, subscribe, destroy }
  const savedSetState = api.setState;
  // change api setState to below , so that when setState is called , it will call setItem to set storage
  api.setState = (state, replace) => {
    savedSetState(state, replace);
    // ignore the return value of setItem, default return value is undefined
    void setItem();
  };

定义第一种返回值

实际上最简单的逻辑已经完成，定义返回值configResult

const configResult = config(
    (...args) => {
      set(...args);
      void setItem();
    },
    get,
    api
  );

定义拿到本地值并合并的方法hydrate

首先这个方法非常长，第一行先定义了一个变量stateFromStorage用来存储合并后的值让我们来看hydrate具体实现：

首先进行storage的错误判断
其次把hasHydrated 改为false，表示暂时没有完成hytrate
然后执行用户定义的hydrationListeners
定义postRehydrationCallback，当用户传入onRehydrateStorage时，定义这个函数，否则定义一个空
执行return，toThenable函数保证了执行的结果一定是一个Promise，toThenable执行返回了一个Promise，并且返回值时本地的存储数据
在第一个then里，如果用户的版本有更新，那么需要用户提供一个migrate函数，如果没有提供，那么报错，如果没有更新，直接返回本地数据
在第二个then里，参数是本地数据或者migrate过后的本地数据，根据用户提供的merge或者默认merge函数，把本地数据和用户提供数据合并，并且更新state，更新本地数据
在第三个then里，执行postRehydrationCallback，并且拿到本地数据，执行用户注册的finishHydrationListeners

整个hydrate函数做的事情就是拿到本地数据，把本地数据合并，并且在恰当的时机执行用户注册的listener

//below is the hydrate func for get the existing state
  let stateFromStorage;
  const hydrate = () => {
    var _a, _b;
    if (!storage) return;
    hasHydrated = false;
    // if any listeners added by user, call hydrationListeners
    hydrationListeners.forEach((cb) => {
      var _a2;
      return cb((_a2 = get()) != null ? _a2 : configResult);
    });
    // if onRehydrateStorage is provided, call it
    const postRehydrationCallback =
      ((_b = options.onRehydrateStorage) == null
        ? void 0
        : _b.call(options, (_a = get()) != null ? _a : configResult)) || void 0;
    return toThenable(storage.getItem.bind(storage))(options.name)
      .then((deserializedStorageValue) => {
        if (deserializedStorageValue) {
          // if version is provided and version is not equal to the version of storage, call migrate func
          if (
            typeof deserializedStorageValue.version === "number" &&
            deserializedStorageValue.version !== options.version
          ) {
            // if migrate func is provided, call it
            if (options.migrate) {
              return options.migrate(
                deserializedStorageValue.state,
                deserializedStorageValue.version
              );
            }
            console.error(
              `State loaded from storage couldn't be migrated since no migrate function was provided`
            );
          } else {
            return deserializedStorageValue.state;
          }
        }
      })
      .then((migratedState) => {
        var _a2;
        stateFromStorage = options.merge(
          migratedState,
          (_a2 = get()) != null ? _a2 : configResult
        );
        set(stateFromStorage, true);
        // operation done , update state and set storage
        return setItem();
      })
      .then(() => {
        // execute post operation of rehydration
        postRehydrationCallback == null
          ? void 0
          : postRehydrationCallback(stateFromStorage, void 0);
        stateFromStorage = get();
        hasHydrated = true;
        finishHydrationListeners.forEach((cb) => cb(stateFromStorage));
      })
      .catch((e) => {
        postRehydrationCallback == null
          ? void 0
          : postRehydrationCallback(void 0, e);
      });
  };

在api对象挂载persist方法

这里挂载了persist的内置方法在api上

  api.persist = {
    setOptions: (newOptions) => {
      options = {
        ...options,
        ...newOptions,
      };
      if (newOptions.storage) {
        storage = newOptions.storage;
      }
    },
    clearStorage: () => {
      storage == null ? void 0 : storage.removeItem(options.name);
    },
    getOptions: () => options,
    rehydrate: () => hydrate(),
    hasHydrated: () => hasHydrated,
    onHydrate: (cb) => {
      hydrationListeners.add(cb);
      return () => {
        hydrationListeners.delete(cb);
      };
    },
    onFinishHydration: (cb) => {
      finishHydrationListeners.add(cb);
      return () => {
        finishHydrationListeners.delete(cb);
      };
    },
  };

persist返回

最后根据用户是否选择跳过hydrate执行hydrate 如果执行了hydrate，那么返回stateFromStorage，否则stateFromStorage为空，返回configResult

  if (!options.skipHydration) {
    hydrate();
  }
  // if configResult exists , return it , otherwise return configResult
  return stateFromStorage || configResult;

总结

整个persist中间件非常清晰，主要是两种情况：

如果用户跳过hydrate，那么返回第一个定义的结果configResult
1. 主要任务是劫持setState，在更新state的同时写入本地
如果用户没有跳过hydrate，那么返回stateFromStorage
1. 主要任务是拿到本地的状态，并且合并
2. 在恰当的时机执行用户注册的listener

使用 Nix 管理及同步 dotfiles Zustand源码解读（一）