【pinia源码】二、defineStore源码解析

前言

【pinia源码】系列文章主要分析pinia的实现原理。该系列文章源码参考pinia v2.0.14

源码地址:https://github.com/vuejs/pinia

官方文档:https://pinia.vuejs.org

本篇文章将分析defineStore的实现。

使用

通过defineStore定义一个store

const useUserStore = defineStore('counter', {
  state: () => ({
    count: 0
  }),
  actions: {
    increment() {
      this.count++
    }
  }
})

// or
const useUserStore = defineStore({
  id: 'counter',
  state: () => ({
    count: 0
  }),
  actions: {
    increment() {
      this.count++
    }
  }
})

// or
const useUserStore = defineStore('counter', () => {
  const count = ref(0)
  
  function increment() {
    count.value++
  }
  return { count, increment }
})

defineStore

export function defineStore(
  idOrOptions: any,
  setup?: any,
  setupOptions?: any
): StoreDefinition {
  let id: string
  let options:
    | DefineStoreOptions<
        string,
        StateTree,
        _GettersTree,
        _ActionsTree
      >
    | DefineSetupStoreOptions<
        string,
        StateTree,
        _GettersTree,
        _ActionsTree
      >

  const isSetupStore = typeof setup === 'function'
  if (typeof idOrOptions === 'string') {
    id = idOrOptions
    options = isSetupStore ? setupOptions : setup
  } else {
    options = idOrOptions
    id = idOrOptions.id
  }

  function useStore(pinia?: Pinia | null, hot?: StoreGeneric): StoreGeneric { // ... }

  useStore.$id = id

  return useStore
}

defineStore函数可以接收三个参数:idOrOptionssetupsetOptions,后两个参数为可选参数。下面是三个defineStore的函数类型定义。

export function defineStore<
  Id extends string,
  S extends StateTree = {},
  G extends _GettersTree = {},
  A /* extends ActionsTree */ = {}
>(
  id: Id,
  options: Omit, 'id'>
): StoreDefinition

export function defineStore<
  Id extends string,
  S extends StateTree = {},
  G extends _GettersTree = {},
  A /* extends ActionsTree */ = {}
  >(options: DefineStoreOptions): StoreDefinition

export function defineStore(
  id: Id,
  storeSetup: () => SS,
  options?: DefineSetupStoreOptions<
    Id,
    _ExtractStateFromSetupStore,
    _ExtractGettersFromSetupStore,
    _ExtractActionsFromSetupStore
    >
): StoreDefinition<
  Id,
  _ExtractStateFromSetupStore,
  _ExtractGettersFromSetupStore,
  _ExtractActionsFromSetupStore
  >

首先在defineStore中声明了三个变量:idoptionsisSetupStore,其中id为定义的store的唯一idoptions为定义store时的optionsisSetupStore代表传入的setup是不是个函数。

然后根据传入的idOrOptions的类型,为idotions赋值。紧接着声明了一个useStore函数,并将id赋给它,然后将其return。截止到此,我们知道defineStore会返回一个函数,那么这个函数具体是做什么的呢?我们继续看useStore的实现。

useStore

function useStore(pinia?: Pinia | null, hot?: StoreGeneric): StoreGeneric {
  // 获取当前实例
  const currentInstance = getCurrentInstance()
  // 测试环境下,忽略提供的参数,因为总是能使用getActivePinia()获取pinia实例
  // 非测试环境下,如果未传入pinia,则会从组件中使用inject获取pinia
  pinia =
    (__TEST__ && activePinia && activePinia._testing ? null : pinia) ||
    (currentInstance && inject(piniaSymbol))
  // 设置激活的pinia
  if (pinia) setActivePinia(pinia)

  // 如果没有activePinia,那么可能没有install pinia,开发环境下进行提示
  if (__DEV__ && !activePinia) {
    throw new Error(
      `[]: getActivePinia was called with no active Pinia. Did you forget to install pinia?\n` +
        `\tconst pinia = createPinia()\n` +
        `\tapp.use(pinia)\n` +
        `This will fail in production.`
    )
  }

  // 设置pinia为激活的pinia
  pinia = activePinia!

  // 从pina._s中查找id否注册过,如果没有被注册,创建一个store并注册在pinia._s中
  if (!pinia._s.has(id)) {
    if (isSetupStore) {
      createSetupStore(id, setup, options, pinia)
    } else {
      createOptionsStore(id, options as any, pinia)
    }

    if (__DEV__) {
      useStore._pinia = pinia
    }
  }

  // 从pinia._s中获取id对应的store
  const store: StoreGeneric = pinia._s.get(id)!

  if (__DEV__ && hot) {
    const hotId = '__hot:' + id
    const newStore = isSetupStore
      ? createSetupStore(hotId, setup, options, pinia, true)
      : createOptionsStore(hotId, assign({}, options) as any, pinia, true)

    hot._hotUpdate(newStore)

    // cleanup the state properties and the store from the cache
    delete pinia.state.value[hotId]
    pinia._s.delete(hotId)
  }

  if (
    __DEV__ &&
    IS_CLIENT &&
    currentInstance &&
    currentInstance.proxy &&
    !hot
  ) {
    const vm = currentInstance.proxy
    const cache = '_pStores' in vm ? vm._pStores! : (vm._pStores = {})
    cache[id] = store
  }

  // 返回store
  return store as any
}

useStore接收两个可选参数:piniahotpinia是个Pinia的实例,而hot只在开发环境下有用,它与模块的热更新有关。

useStore中会首先获取当前组件实例,如果存在组件实例,使用inject(piniaSymbol)获取pinia(在install中会进行provide),并将其设置为activePinia,然后在activePinia._s中查找是否有被注册为idstore,如果没有则创建store,将其注册到activePinia._s中。最后返回activePinia._sid对应的store

现在我们知道useStore函数,最终会返回一个store。那么这个store是什么呢?它是如何创建的呢?在useStore中根据不同情况中有两中方式来创建store,分别是:createSetupStorecreateOptionsStore。这两个方式的使用条件是:如果defineStore第二个参数是个function调用createSetupStore,相反调用createOptionsStore

createSetupStore

createSetupStore函数代码过长,这里就不贴完整代码了。createSetupStore可接收参数如下:

参数 说明
$id 定义storeid
setup 一个可以返回state的函数
options defineStoreoptions
pinia Pinia实例
hot 是否启用热更新 可选
isOptionsStore 是否使用options声明的store 可选

createSetupStore代码有500多行,如果从头开始看的话,不容易理解。我们可以根据createSetupStore的用途,从其核心开始看。因为createSetupStore是需要创建store,并将store注册到pinia._s中,所以createSetupStore中可能需要创建store,我们找到创建store的地方。

const partialStore = {
  _p: pinia,
  // _s: scope,
  $id,
  $onAction: addSubscription.bind(null, actionSubscriptions),
  $patch,
  $reset,
  $subscribe(callback, options = {}) {
    const removeSubscription = addSubscription(
      subscriptions,
      callback,
      options.detached,
      () => stopWatcher()
    )
    const stopWatcher = scope.run(() =>
      watch(
        () => pinia.state.value[$id] as UnwrapRef,
        (state) => {
          if (options.flush === 'sync' ? isSyncListening : isListening) {
            callback(
              {
                storeId: $id,
                type: MutationType.direct,
                events: debuggerEvents as DebuggerEvent,
              },
              state
            )
          }
        },
        assign({}, $subscribeOptions, options)
      )
    )!

    return removeSubscription
  },
  $dispose,
} as _StoreWithState

if (isVue2) {
  partialStore._r = false
}

const store: Store = reactive(
  assign(
    __DEV__ && IS_CLIENT
      ? // devtools custom properties
        {
          _customProperties: markRaw(new Set()),
          _hmrPayload,
        }
      : {},
    partialStore
  )
) as unknown as Store

pinia._s.set($id, store)

store是用reactive包装的一个响应式对象,reactive所包装的对象是由partialStore通过Object.assign进行复制的。partialStore中定义了很多方法,这些方法都是暴露给用户操作store的一些接口,如$onAction可设置actions的回调、$patch可更新store中的state$dispose可销毁store

在调用完pinia._s.set($id, store)之后,会执行setup,获取所有的数据。setup的执行会在创建pinia实例时创建的effectScope中运行,而且会再单独创建一个effectScope,用来单独执行setup.

const setupStore = pinia._e.run(() => {
  scope = effectScope()
  return scope.run(() => setup())
})!

然后遍历setupStore的属性:如果propkey对应的值)为ref(不为computed)或reactive,则将keyprop同步到pina.state.value[$id]中;如果propfunction,则会使用wrapAction包装prop,并将包装后的方法赋值给setupStore[key],以覆盖之前的值,同时将包装后的方法存入optionsForPlugin.actions中。

for (const key in setupStore) {
  const prop = setupStore[key]

  // 如果prop是ref(但不是computed)或reactive
  if ((isRef(prop) && !isComputed(prop)) || isReactive(prop)) {
    if (__DEV__ && hot) {
      set(hotState.value, key, toRef(setupStore as any, key))
    } else if (!isOptionsStore) {
      if (initialState && shouldHydrate(prop)) {
        if (isRef(prop)) {
          prop.value = initialState[key]
        } else {
          mergeReactiveObjects(prop, initialState[key])
        }
      }

      // 将对应属性同步至pinia.state中
      if (isVue2) {
        set(pinia.state.value[$id], key, prop)
      } else {
        pinia.state.value[$id][key] = prop
      }
    }

    if (__DEV__) {
      _hmrPayload.state.push(key)
    }
  } else if (typeof prop === 'function') { // 如果prop是function
    // 使用wrapAction包装prop,在wrapAction会处理afeterCallback、errorCallback
    const actionValue = __DEV__ && hot ? prop : wrapAction(key, prop)

    // 将actionsValue添加到setupStore中,覆盖原来的function
    if (isVue2) {
      set(setupStore, key, actionValue)
    } else {
      setupStore[key] = actionValue
    }

    if (__DEV__) {
      _hmrPayload.actions[key] = prop
    }

    // 将function类型的prop存入optionsForPlugin.actions中
    optionsForPlugin.actions[key] = prop
  } else if (__DEV__) {
    if (isComputed(prop)) {
      _hmrPayload.getters[key] = isOptionsStore
        ? // @ts-expect-error
        options.getters[key]
        : prop
      if (IS_CLIENT) {
        const getters: string[] =
          setupStore._getters || (setupStore._getters = markRaw([]))
        getters.push(key)
      }
    }
  }
}

接下来我们看下wrapAction是如何进行包装function类型上的prop

function wrapAction(name: string, action: _Method) {
  return function (this: any) {
    setActivePinia(pinia)
    const args = Array.from(arguments)

    const afterCallbackList: Array<(resolvedReturn: any) => any> = []
    const onErrorCallbackList: Array<(error: unknown) => unknown> = []
    function after(callback: _ArrayType) {
      afterCallbackList.push(callback)
    }
    function onError(callback: _ArrayType) {
      onErrorCallbackList.push(callback)
    }

    triggerSubscriptions(actionSubscriptions, {
      args,
      name,
      store,
      after,
      onError,
    })

    let ret: any
    try {
      ret = action.apply(this && this.$id === $id ? this : store, args)
    } catch (error) {
      triggerSubscriptions(onErrorCallbackList, error)
      throw error
    }

    // 如果结果是promise,在promise中触发afterCallbackList及onErrorCallbackList
    if (ret instanceof Promise) {
      return ret
        .then((value) => {
          triggerSubscriptions(afterCallbackList, value)
          return value
        })
        .catch((error) => {
          triggerSubscriptions(onErrorCallbackList, error)
          return Promise.reject(error)
        })
    }

    triggerSubscriptions(afterCallbackList, ret)
    return ret
  }
}

wrapAction首先返回一个函数,在这个函数中,首先将pinia设置为activePinia,触发actionSubscriptions中的函数,然后执行action函数,如果执行过程中出错,会执行onErrorCallbackList中的errorCallback,如果没有出错的话,执行afterCallbackList中的afterCallback,最后将action的返回结果return

wrapAction中的actionSubscriptions是个什么呢?

其实actionSubscriptions中的callback就是是通过store.$onAction添加的回调函数;在执行actionSubscriptions中的callback过程中,会将对应callback添加到afterCallbackListonErrorCallbackList中。例如:

store.$onAction(({ after, onError, name, store }) => {
  after((value) => {
    console.log(value)
  })
  
  onError((error) => {
    console.log(error)
  })
})

遍历完setupStore之后,会将setupStore合并至storestore的原始对对象中,以方便使用storeToRefs()检索响应式对象。

if (isVue2) {
  Object.keys(setupStore).forEach((key) => {
    set(
      store,
      key,
      setupStore[key]
    )
  })
} else {
  assign(store, setupStore)
  assign(toRaw(store), setupStore)
}

紧接着拦截store.$stategetset方法:当调用store.$state时,能够从pinia.state.value找到对应的state;当使用store.$state = xxx去修改值时,则调用$patch方法修改值。

Object.defineProperty(store, '$state', {
  get: () => (__DEV__ && hot ? hotState.value : pinia.state.value[$id]),
  set: (state) => {
    /* istanbul ignore if */
    if (__DEV__ && hot) {
      throw new Error('cannot set hotState')
    }
    $patch(($state) => {
      assign($state, state)
    })
  },
})

截止到此,store就准备完毕。如果在Vue2环境下,会将store._r设置为true。

if (isVue2) {
  store._r = true
}

接下来就需要调用使用use方法注册的plugins

pinia._p.forEach((extender) => {
  if (__DEV__ && IS_CLIENT) {
    const extensions = scope.run(() =>
      extender({
        store,
        app: pinia._a,
        pinia,
        options: optionsForPlugin,
      })
    )!
    Object.keys(extensions || {}).forEach((key) =>
      store._customProperties.add(key)
    )
    assign(store, extensions)
  } else {
    // 将plugin的结果合并到store中
    assign(
      store,
      scope.run(() =>
        extender({
          store,
          app: pinia._a,
          pinia,
          options: optionsForPlugin,
        })
      )!
    )
  }
})

最后返回store

if (
  initialState &&
  isOptionsStore &&
  (options as DefineStoreOptions).hydrate
) {
  ;(options as DefineStoreOptions).hydrate!(
    store.$state,
    initialState
  )
}

isListening = true
isSyncListening = true
return store

接下来看下store中的几个方法:

$onAction

在每个action中添加回调函数。回调接收一个对象参数:该对象包含nameactionkey值)、store(当前store)、after(添加action执行完之后的回调)、onError(添加action执行过程中的错误回调)、argsaction的参数)属性。

示例:

// 统计add action的调用次数
let count = 0, successCount = 0, failCount = 0
store.$onAction(({ name, after, onError }) => {
  if (name === 'add') {
    count++
    after((resolveValue) => {
      successCount++
      console.log(resolveValue)
    })
  
    onError((error) => {
      failCount++
      console.log(error)
    })
  }
})

$onAction内部通过发布订阅模式实现。在pinia中有个专门的订阅模块subscriptions.ts,其中包含两个主要方法:addSubscription(添加订阅)、triggerSubscriptions(触发订阅)。

addSubscription可接收四个参数:subscriptions(订阅列表)、callback(添加的订阅函数)、detached(游离的订阅,如果为false在组件卸载后,自动移除订阅;如果为true,不会自动移除订阅)、onCleanup(订阅被移除时的回调)

triggerSubscriptions接收两个参数:subscriptions(订阅列表)、argsaction的参数列表)

export function addSubscription(
  subscriptions: T[],
  callback: T,
  detached?: boolean,
  onCleanup: () => void = noop
) {
  subscriptions.push(callback)

  const removeSubscription = () => {
    const idx = subscriptions.indexOf(callback)
    if (idx > -1) {
      subscriptions.splice(idx, 1)
      onCleanup()
    }
  }

  if (!detached && getCurrentInstance()) {
    onUnmounted(removeSubscription)
  }

  return removeSubscription
}

export function triggerSubscriptions(
  subscriptions: T[],
  ...args: Parameters
) {
  subscriptions.slice().forEach((callback) => {
    callback(...args)
  })
}

$onAction通过addSubscription.bind(null, actionSubscriptions)实现。

如何触发订阅?

首先在store的初始化过程中,会将action使用wrapAction函数进行包装,wrapAction返回一个函数,在这个函数中会先触发actionSubscriptions,这个触发过程中会将afterCallbackonErrorCallback添加到对应列表。然后调用action,如果调用过程中出错,则触发onErrorCallbackList,否则触发afterCallbackList。如果action的结果是Promise的话,则在then中触发onErrorCallbackList,在catch中触发onErrorCallbackList。然后会将包装后的action覆盖原始action,这样每次调用action时就是调用的包装后的action

$patch

使用$patch可以更新state的值,可进行批量更新。$patch接收一个partialStateOrMutator参数,它可以是个对象也可以是个方法。

示例:

store.$patch((state) => {
  state.name = 'xxx'
  state.age = 14
})
// or
store.$patch({
  name: 'xxx',
  age: 14
})

$patch源码:

function $patch(
  partialStateOrMutator:
    | _DeepPartial>
    | ((state: UnwrapRef) => void)
): void {
  // 合并的相关信息
  let subscriptionMutation: SubscriptionCallbackMutation
  // 是否触发状态修改后的回调,isListening代表异步触发,isSyncListening代表同步触发
  // 此处先关闭回调的触发,防止修改state的过程中频繁触发回调
  isListening = isSyncListening = false
  if (__DEV__) {
    debuggerEvents = []
  }
  // 如果partialStateOrMutator是个function,执行方法,传入当前的store
  if (typeof partialStateOrMutator === 'function') {
    partialStateOrMutator(pinia.state.value[$id] as UnwrapRef)
    subscriptionMutation = {
      type: MutationType.patchFunction,
      storeId: $id,
      events: debuggerEvents as DebuggerEvent[],
    }
  } else { // 如果不是function,则调用mergeReactiveObjects合并state
    mergeReactiveObjects(pinia.state.value[$id], partialStateOrMutator)
    subscriptionMutation = {
      type: MutationType.patchObject,
      payload: partialStateOrMutator,
      storeId: $id,
      events: debuggerEvents as DebuggerEvent[],
    }
  }
  // 当合并完之后,将isListening、isSyncListening设置为true,意味着可以触发状态改变后的回调函数了
  const myListenerId = (activeListener = Symbol())
  nextTick().then(() => {
    if (activeListener === myListenerId) {
      isListening = true
    }
  })
  isSyncListening = true
  // 因为在修改pinia.state.value[$id]的过程中关闭(isSyncListening与isListening)了监听,所以需要手动触发订阅列表
  triggerSubscriptions(
    subscriptions,
    subscriptionMutation,
    pinia.state.value[$id] as UnwrapRef
  )
}

$reset

通过构建一个新的state objectstate重置为初始状态。只在options配置下生效。如果是setup配置,开发环境下报错。

store.$reset = function $reset() {
  // 重新执行state,获取一个新的state
  const newState = state ? state() : {}
  // 通过$patch,使用assign将newState合并到$state中
  this.$patch(($state) => {
    assign($state, newState)
  })
}

$subscribe

设置state改变后的回调,返回一个移除回调的函数。可接受两个参数:callback(添加的回调函数)、options:{detached, flush, ...watchOptions}detachedaddSubscription中的detachedflush代表是否同步触发回调,可取值:sync)。

示例:

store.$subribe((mutation: {storeId, type, events}, state) => {
  console.log(storeId)
  console.log(type)
  console.log(state)
}, { detached: true, flush: 'sync' })

$subscribe源码:

function $subscribe(callback, options = {}) {
  // 将callback添加到subscriptions中,以便使用$patch更新状态时,触发回调
  // 当使用removeSubscription移除callback时,停止对pinia.state.value[$id]监听
  const removeSubscription = addSubscription(
    subscriptions,
    callback,
    options.detached,
    () => stopWatcher()
  )
  const stopWatcher = scope.run(() =>
    // 监听pinia.state.value[$id],以触发callback,当使用$patch更新state时,不会进入触发这里的callback
    watch(
      () => pinia.state.value[$id] as UnwrapRef,
      (state) => {
        if (options.flush === 'sync' ? isSyncListening : isListening) {
          callback(
            {
              storeId: $id,
              type: MutationType.direct,
              events: debuggerEvents as DebuggerEvent,
            },
            state
          )
        }
      },
      assign({}, $subscribeOptions, options)
    )
  )!

  return removeSubscription
}

callback中的第一个参数中有个type属性,表示是通过什么方式更新的state,它有三个值:

  1. MutationType.direct:通过state.name='xxx'/store.$state.name='xxx'等方式修改
  2. MutationType.patchObject:通过store.$patch({ name: 'xxx' })方式修改
  3. MutationType.patchFunction:通过store.$patch((state) => state.name='xxx')方式修改

$dispose

销毁store

function $dispose() {
  // 停止监听
  scope.stop()
  // 清空subscriptions及actionSubscriptions
  subscriptions = []
  actionSubscriptions = []
  // 从pinia._s中删除store
  pinia._s.delete($id)
}

createOptionsStore

createOptionsStore可接收参数如下:

参数 说明
id 定义storeid
options defineStoreoptions
pinia Pinia实例
hot 是否启用热更新 可选
function createOptionsStore<
  Id extends string,
  S extends StateTree,
  G extends _GettersTree,
  A extends _ActionsTree
>(
  id: Id,
  options: DefineStoreOptions,
  pinia: Pinia,
  hot?: boolean
): Store {
  const { state, actions, getters } = options

  const initialState: StateTree | undefined = pinia.state.value[id]

  let store: Store

  function setup() {
    // 如果pinia.state.value[id]不存在,进行初始化
    if (!initialState && (!__DEV__ || !hot)) {
      if (isVue2) {
        set(pinia.state.value, id, state ? state() : {})
      } else {
        pinia.state.value[id] = state ? state() : {}
      }
    }

    // 将pinia.state.value[id]各属性值转为响应式对象
    const localState =
      __DEV__ && hot
        ? // use ref() to unwrap refs inside state TODO: check if this is still necessary
          toRefs(ref(state ? state() : {}).value)
        : toRefs(pinia.state.value[id])

    // 处理getters,并将处理后的getters和actions合并到localState中
    return assign(
      localState,
      actions,
      Object.keys(getters || {}).reduce((computedGetters, name) => {
        computedGetters[name] = markRaw(
          computed(() => {
            setActivePinia(pinia)
            const store = pinia._s.get(id)!
            
            if (isVue2 && !store._r) return

            return getters![name].call(store, store)
          })
        )
        return computedGetters
      }, {} as Record)
    )
  }

  // 利用createSetupStore创建store
  store = createSetupStore(id, setup, options, pinia, hot, true)

  // 重写store.$reset
  store.$reset = function $reset() {
    const newState = state ? state() : {}
    this.$patch(($state) => {
      assign($state, newState)
    })
  }

  return store as any
}

createOptionsStore中会根据传入参数构造一个setup函数,然后通过createSetupStore创建一个store,并重写store.$reset方法,最后返回store

这个setup函数中会将state()的返回值赋值给pinia.state.value[id],然后将pinia.state.value[id]进行toRefs,得到localState,最后将处理后的gettersactions都合并到localState中,将其返回。对于getters的处理:将每个getter函数都转成一个计算属性。

总结

defineStore返回一个useStore函数,通过执行useStore可以获取对应的store。调用useStore时我们并没有传入id,为什么能准确获取store呢?这是因为useStore是个闭包,在执行useStore执行过程中会自动获取id

获取store的过程:

  1. 首先获取组件实例
  2. 使用inject(piniaSymbol)获取pinia实例
  3. 判断pinia._s中是否有对应id的键,如果有直接取对应的值作为store,如果没有则创建store

store创建流程分两种:setup方式与options方式

setup方式:

  1. 首先在pinia.state.value中添加键为$id的空对象,以便后续赋值
  2. 使用reactive声明一个响应式对象store
  3. store存至pinia._s
  4. 执行setup获取返回值setupStore
  5. 遍历setupStore的键值,如果值是ref(不是computed)或reactive,将键值添加到pinia.state.value[$id]中;如果值时function,首先将值使用wrapAction包装,然后用包装后的function替换setupStore中对应的值
  6. setupStore合并到store
  7. 拦截store.$state,使get操作可以正确获取pinia.state.value[$id]set操作使用this.$patch更新
  8. 调用pinia._p中的扩展函数,扩展store

options方式:

  1. options中提取stategetteractions
  2. 构建setup函数,在setup函数中会将getter处理成计算属性
  3. 使用setup方式创建store
  4. 重写store.$reset

你可能感兴趣的:(【pinia源码】二、defineStore源码解析)