pinia源码分析 - 技术栈

本文同步在个人博客shymean.com上，欢迎关注

在项目中使用Vue3和Pinia已经很长时间了，Pinia作为一个组合式Api的Vue状态管理库，在使用过程中也有一些心得。

本文将从源码的角度分析Pinia，最后会实现一个简易版的Pinia。

参考

Pinia 官网

本文完整代码已放在github上面了。

Vue3的全局状态

全局reactive对象

事实上，借助于Vue3的响应式系统，我们不使用任何状态管理库也可以实现全局状态，以及跨组件数据共享。

首先通过reactive定义全局对象

ts 复制代码

import {reactive} from "vue";

export const globalState = reactive({
  count:1
})
export function increase(){
  globalState.count +=1
}
export function decrease(){
  globalState.count -=1
}

然后，只需要在对应的组件内引入globalState，就可以使用上面的globalState.count状态了。

当修改数据之后，所有使用globalState.count的组件都会同步更新。

这种写法的缺点在于如果整个项目使用了SSR，全局变量无法实现状态隔离。

Pinia的使用

关于pinia的使用，可以参考官方文档，这里只是先回顾一下，以及找到源码调试的入口

ts 复制代码

export const useCounter = defineStore('counter', {
	state:()=>{
		return {
  		count: 1
		}
	},
	getters: {
		
	},
	actions: {
		
	}
})

在初始化Vue应用之后

ts 复制代码

const app = createApp(App)
app.use(createPinia()).use(xxx)

在需要全局状态的组件中

ts 复制代码

const counterStore = useCounter()
console.log(counterStore.count)

核心源码

接下来看看Pinia的源码

createPinia

ts 复制代码

export function createPinia(): Pinia {
  const scope = effectScope(true)
  const state = scope.run<Ref<Record<string, StateTree>>>(() =>
    ref<Record<string, StateTree>>({})
  )!
	const pinia: Pinia = markRaw({
    install(app: App) {
      // 应用注册的时候调用
      setActivePinia(pinia)
    },
    state, // 全局状态
    _e: scope,
    _s: new Map<string, StoreGeneric>(), // 保存defineStore定义的store，使用store id作为键值
  })
  return pinia
}

这个setActivePinia就是更新一下全局的activePinia变量

ts 复制代码

export let activePinia: Pinia | undefined
export const setActivePinia: _SetActivePinia = (pinia) => (activePinia = pinia)

可以看出，pinia对象就是用来保存一下全局状态而已。

在单页应用中，一般只会有一次createPinia的调用，在SSR中，需要为每一次请求生成一个独立的pinia对象，用于隔离状态。

effectScope

参考：effectScope RFC

在常规的组件setup中，vue会收集属性的依赖并绑定到组件上，这样当属性变化时就可以触发组件的render等；当组件被卸载时，相关的依赖会被自动释放。

但在组件之外或作为独立包时，要达到类似的效果，就可以使用effectScope。

ts 复制代码

import { effectScope, reactive, watch } from 'vue';

const scope = effectScope();

const state = reactive({
  count: 0,
})

scope.run(() => {
  watch(()=>state.count,() => {
    console.log('Count has changed:', state.count);
  });

  state.count += 1 // 会触发watch
});

// 模拟组件onUnMounted
setTimeout(()=>{
  console.log('scope stop')
  scope.stop();
  state.count++; // 不会再触发 watch 的回调
})

作用域可以嵌套，因此effectScope接收一个detached参数，表示嵌套的子作用域是否要跟外层的副作用域独立开来。

detached默认为false，表示不独立，这样当父作用域stop之后，子作用域也会stop；如果设置为true，则即使副作用域stop之后，子作用域也不会停止，直到他调用了自己的stop方法。

ts 复制代码

let nestedScope
scope.run(() => {
  watch(()=>state.count,() => {
    console.log('Count has changed:', state.count);
  });

  nestedScope = effectScope(true) // 独立
  nestedScope.run(() => {
    watch(()=>state.count,() => {
      console.log('Count has changed in nested:', state.count);
    });
  })

  state.count += 1 // 会触发watch
});


setTimeout(()=>{
  console.log('scope stop')
  scope.stop();
  state.count++; // scope不会再触发 watch 的回调，但nestedScope内的watch还会触发
})

defineStore

先看常规的setup形式定义的store

ts 复制代码

function defineStore(options){
    function useStore(pinia?: Pinia){
        // 可以传入自定义的pinia，在SSR等场景下可以根据请求上下文区分各自的全局状态
        if (pinia) setActivePinia(pinia)
        pinia = activePinia!
        // 当第一次调用useStore的时候才初始化
        if (!pinia._s.has(id)) {
          createSetupStore(id, setup, options as any, pinia)
        }	
	const store = pinia._s.get(id)!
        return  store
    }
    return useStore
}

createSetupStore这个方法很长，但他的核心代码很简单

ts 复制代码

function createSetupStore(){
  const store = {
    _p: pinia,
    // _s: scope,
    $id,
    // ... 其他工具方法如$patch、$dispose等
  }
  const setupStore = pinia._e.run(() => (scope = effectScope()).run(setup)!)
  assign(store, setupStore)
  return store
}

注意上面这里的effectScope没有传入detached，这样可以与父作用域（也就是在createPinia那里创建的effectScope(true)）共享生命周期。

最后，就可以通过useXXStore获取到全局的状态和方法（这也是在使用pinia中最常用的方法）。

ts 复制代码

const store = useXXStore()
console.log(store.xxState)

至此，Pinia的核心源码好像就看完了~（捂脸

类型推断

Pinia还有一个非常重要的特性：类型推断实现的智能提示。

在Vuex中，最大的问题是所有actions、mutation等都需要字符串

js 复制代码

this.$store.commit("user/setUserInfo",{})
this.$store.commit("user/fetchUserInfo", {id:xxx})
console.log(this.$store.state.user.UserInfo)

虽然WebStrom支持通过字符串跳转到源码处，但这种手写字符串的做法还是有很大的心智负担。

借助TS，Pinia实现了智能提示，接下来我们需要学习一下Pinia源码中的类型实现。

defineStore返回值

在useXXStore之后，得到的是一个store对象，如何推断出他上面的state类型呢？

ts 复制代码

export type StateTree = Record<string | number | symbol, any>

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

看一下defineStore返回值的类型StoreDefinition

ts 复制代码

export interface StoreDefinition<
  Id extends string = string,
  S extends StateTree = StateTree,
  G /* extends GettersTree<S>*/ = _GettersTree<S>,
  A /* extends ActionsTree */ = _ActionsTree
> {
  (pinia?: Pinia | null | undefined, hot?: StoreGeneric): Store<Id, S, G, A>
  $id: Id
  _pinia?: Pinia
}

state函数返回值推断

从StoreDefinition的定义可以看到useXXStore的返回值类型 Store<Id, S, G, A>，每个泛型

S，state对应_ExtractStateFromSetupStore<SS>
G，getter对应_ExtractGettersFromSetupStore<SS>
A，action对应_ExtractActionsFromSetupStore<SS>

SS就是defineStore传入的setup函数的返回值。

我们来分析一下Stated的类型_ExtractStateFromSetupStore是如何工作的

ts 复制代码

export type _UnwrapAll<SS> = { [K in keyof SS]: UnwrapRef<SS[K]> }
export type _ExtractStateFromSetupStore_Keys<SS> = keyof {
  [K in keyof SS as SS[K] extends _Method | ComputedRef ? never : K]: any
}
export type _ExtractStateFromSetupStore<SS> = SS extends undefined | void
  ? {}
  : _ExtractStateFromSetupStore_Keys<SS> extends keyof SS // 这一步判断感觉是多余的了
  ? _UnwrapAll<Pick<SS, _ExtractStateFromSetupStore_Keys<SS>>>
  : never

_ExtractStateFromSetupStore_Keys主要用来排除SS类型中的计算属性和方法，得到一个由state名字组成的字符串字面量联合类型

然后Pick<SS, _ExtractStateFromSetupStore_Keys<SS>>返回就可以返回 SS 中排除了计算属性和方法的其余属性类型

最后，通过_UnwrapAll返回数据的原始类型，其内部使用了UnwrapRef，这是Vue3的一个工具函数

比如 Ref<number>，经过UnwrapRef<Ref<number>> 将返回 number，这保证了最终返回的类型是 SS 中符合条件的属性的实际类型，而不是被 Ref 或其他包装类型包裹的类型

我们可以直接在源码这里修改一下进行测试

可以看见最终StateTest上面得到到字段 "a" | "a2"。

ts 复制代码

type StateTest = _ExtractStateFromSetupStore<{
  a: Ref<number>,
  a2: Ref<boolean>,
  b: ComputedRef<number>,
  c: Function
}>

type a = StateTest['a'] // number
type a2 = StateTest['a2'] // boolean

这样就可以将setup函数返回值里面所有状态都提取出来，然后就可以直到store上面的到底有哪些类型了。

同样地，Getter就是从SS中提取出类型为ComputedRef的属性

ts 复制代码

export type _ExtractGettersFromSetupStore_Keys<SS> = keyof {
  [K in keyof SS as SS[K] extends ComputedRef ? K : never]: any
}

而Action就是从SS中提取出类型为函数的属性

ts 复制代码

export type _Method = (...args: any[]) => any
export type _ExtractActionsFromSetupStore_Keys<SS> = keyof {
  [K in keyof SS as SS[K] extends _Method ? K : never]: any
}

Store

现在我们了解了S、G、A的泛型类型，再来看看Store的定义就可以了

完整的Store类型是

ts 复制代码

export type Store<
  Id extends string = string,
  S extends StateTree = {},``
  G /* extends GettersTree<S>*/ = {},
  // has the actions without the context (this) for typings
  A /* extends ActionsTree */ = {}
> = _StoreWithState<Id, S, G, A> &
  UnwrapRef<S> &
  _StoreWithGetters<G> &
  // StoreWithActions<A> &
  (_ActionsTree extends A ? {} : A) &
  PiniaCustomProperties<Id, S, G, A> &
  PiniaCustomStateProperties<S>

其中

_StoreWithState定义了Store的一些原始属性方法，诸如$id、$patch、$reset等，这里不展示具体代码

UnwrapRef<S>这里就是将state的属性合并到Store上面

_ActionsTree就是将setup返回的方法合并到Store上面

对于_StoreWithGetters，则是对从setup返回的getter合并到Store上面

typescript 复制代码

export type _StoreWithGetters<G> = {
  readonly [k in keyof G]: G[k] extends (...args: any[]) => infer R
    ? R
    : UnwrapRef<G[k]>
}

[k in keyof G] : 使用映射类型，遍历 G 类型中的所有键 k。
G[k] extends (...args: any[]) => infer R ? R : UnwrapRef<G[k]>:
- 对于 G 中的每个键 k，检查 G[k] 的类型。
- 如果 G[k] 的类型是一个函数，即函数类型，表示它是一个 getter 函数，则返回该函数的返回类型 infer R，目前看起来，我好像没有见过有G[k] extends (...args: any[]) => infer R这种类型的计算属性，在上面_ExtractGettersFromSetupStore_Keys这里提取的也是ComputedRef形式的键值出来，因此这个三元运算符的必要性还需要研究一下。
- 如果 G[k] 不是函数，即不是 getter 函数，使用 UnwrapRef 将其类型解包，确保得到的是实际的值而不是 Ref 对象。
readonly: 将对象设置为只读，确保 Getters 在 Store 中是只读的。

mini-pinia

根据我们看到的源码，我们可以100来行代码还原一个mini版本的pinia。这个版本

只考虑Vue3
只支持steup形式的store，支持类型推断
不实现$patch、$reset等方法
不支持插件

实现

首先是createPinia

ts 复制代码

import { effectScope, EffectScope, ref, Ref } from 'vue'

export interface Pinia {
  install: Function
  scope: EffectScope
  stores: Record<string, any>
  state: Ref<any>
}

export let activePinia: Pinia | undefined

export function setActivatePinia(p: Pinia) {
  activePinia = p
}
export function createPinia() {
  const scope = effectScope(true)
  const state = scope.run(() => {
    return ref({})
  }) as Ref<any>

  const pinia: Pinia = {
    install() {
      setActivatePinia(pinia)
    },
    scope,
    state,
    stores: {} as Record<string, any>,
  }
  return pinia
}

然后是defineStore，首先需要定义一下类型

首先是Store的类型

ts 复制代码

export type _Method = (...args: any[]) => any
export type _UnwrapAll<SS> = { [K in keyof SS]: UnwrapRef<SS[K]> }

export type _StoreWithGetters<G> = {
  readonly [k in keyof G]: UnwrapRef<G[k]>
}
export type StateTree = Record<string | number | symbol, any>

export type Store<
  Id extends string = string,
  S extends StateTree = {},
  G = {},
  A = {}
> = UnwrapRef<S> &
  _StoreWithGetters<G> &
  A & {
    $id: Id
  }

然后是defineStore的函数签名，通过storeSetup: () => SS保存setup函数返回值的类型，方便后面做类型推断

ts 复制代码

interface StoreDefinition<
  Id extends string = string,
  S extends StateTree = StateTree,
  G = {},
  A = {}
> {
  (pinia?: Pinia | undefined): Store<Id, S, G, A>
}

export function defineStore<Id extends string, SS>(
  id: Id,
  storeSetup: () => SS
): StoreDefinition<
  Id,
  _ExtractStateFromSetupStore<SS>,
  _ExtractGettersFromSetupStore<SS>,
  _ExtractActionsFromSetupStore<SS>
>

接着从SS上面拆分State、Getter和Action

ts 复制代码

export type _ExtractStateFromSetupStore_Keys<SS> = keyof {
  [K in keyof SS as SS[K] extends _Method | ComputedRef ? never : K]: any
}
export type _ExtractStateFromSetupStore<SS> = SS extends undefined | void
  ? {}
  : _UnwrapAll<Pick<SS, _ExtractStateFromSetupStore_Keys<SS>>>

export type _ExtractGettersFromSetupStore_Keys<SS> = keyof {
  [K in keyof SS as SS[K] extends ComputedRef ? K : never]: any
}

export type _ExtractGettersFromSetupStore<SS> = SS extends undefined | void
  ? {}
  : Pick<SS, _ExtractGettersFromSetupStore_Keys<SS>>

export type _ExtractActionsFromSetupStore_Keys<SS> = keyof {
  [K in keyof SS as SS[K] extends _Method ? K : never]: any
}
export type _ExtractActionsFromSetupStore<SS> = SS extends undefined | void
  ? {}
  : Pick<SS, _ExtractActionsFromSetupStore_Keys<SS>>

类型定义完成，再来实现一下defineStore函数

ts 复制代码

export function defineStore(id: string, setup: any) {
  return function useStore(pinia?: Pinia) {
    if (pinia) {
      setActivatePinia(pinia)
    }
    pinia = activePinia
    if (!pinia) {
      throw new Error('no active pinia')
    }
    if (!pinia.stores[id]) {
      createStore(pinia, id, setup)
    }
    return pinia.stores[id] as any
  }
}

function createStore(pinia: Pinia, id: string, setup: any) {
  const store = {}
  const setupStore = pinia.scope.run(setup)
  Object.assign(store, setupStore)
  pinia.stores[id] = store as any
}

最后再来实现工具函数storeToRefs，跟Pinia一样，这里只返回state和getter，至于不响应的数据和action，直接通过store调用即可，这里无需返回

ts 复制代码

type ToComputedRefs<T> = {
  [K in keyof T]: ToRef<T[K]> extends Ref<infer U> ? ComputedRef<U> : ToRef<T[K]>
}

export type StoreToRefs<SS> = ToRefs<_ExtractStateFromSetupStore<SS>> &
  ToComputedRefs<_ExtractGettersFromSetupStore<SS>>

export function storeToRefs<SS extends {}>(store: SS): StoreToRefs<SS> {
  store = toRaw(store)
  const refs = {} as StoreToRefs<SS>
  for (const key in store) {
    const value = store[key]
    if (isRef(value) || isReactive(value)) {
      // @ts-ignore
      refs[key] = toRef(store, key)
    }
  }
  return refs
}

OK，大功告成

使用

这个mini-pinia使用起来，跟常规的pinia没有啥区别

vue 复制代码

<script setup lang="ts">
import HelloWorld from './components/HelloWorld.vue'
import { useCounterStore } from './store/counter.ts'
import { storeToRefs } from './pinia/store.ts'

const counterStore = useCounterStore()
const { count } = storeToRefs(counterStore)
</script>

<template>
  <div>
    <button @click="counterStore.decrement">-</button>
    <button>
      {{ count }}
    </button>
    <button @click="counterStore.increment">+</button>
    <div>
      <div>children:</div>
      <HelloWorld />
    </div>
  </div>
</template>

注意Hello World组件中同样使用了counterStore，可以观察跨组件的状态是否一致，这里不再展开。

mini-pinia完整代码已放在github上面了。

其他功能

虽然mini-pinia的代码虽然很短，但揭露了Pinia的核心原理。可以看出，相较于Vuex，Pinia的确很精简。

上面内容可以忽略了很多细节，比如pinia热更新、插件系统等。

options转setup

definStore有两种写法：option写法和setup写法

由于之前一直使用vuex，我习惯了编写option写法，pinia废弃了mutation写法，只需要编写state、getters和actions就行了。

在实现上，option写法最终最转成setup写法，

ts 复制代码

if (isSetupStore) {
  createSetupStore(id, setup, options, pinia)
} else {
  createOptionsStore(id, options as any, pinia)
}

ts 复制代码

function createOptionsStore(id, options, pinia){
  function setup(){
  	// 根据option创建setup函数
  }
	store = createSetupStore(id, setup, options, pinia, hot, true)
}

由于option写法的state是一个函数，因此内置了一个$reset方法，如果是直接使用setup写法的话，需要自己实现。

setup写法更加灵活，比如某个store依赖其他的store，不需要再每个getter或者action中重复使用useXX，但需要注意的是这种写法不能在两个setup的store中互相引用，导致出现无限循环的情况。

option选项写法的最大问题是逻辑可能会分散，不过如果保证store按照功能拆分的比较写，单个store的状态和方法也不会写的太多。

如果对pinia不是很熟悉的话，可以先根据文档的建议，使用option写法。

插件系统

使用插件可以扩展单个store、添加新属性，包装action等。

在源码实现上，每个pinia对象都有一个_p用来保存插件列表，并在createStore的时候，遍历插件列表，同时将每个插件的返回值合并到store上面

ts 复制代码

pinia._p.forEach((extender) => {
	assign(
    store,
    scope.run(() =>
      extender({
        store: store as Store,
        app: pinia._a,
        pinia,
        options: optionsForPlugin,
      })
    )!
  )
})

在返回值这里，就可以返回同名的方法，实现包装action的效果；或者返回不同名的state，实现扩展state的效果。

一个常见的使用状态管理插件的场景是管理action loading。

比如dva-loading、vuex-loading之类的效果：某个xxAction运行时，设置一个全局状态xxActionLoading，展示加载动画等。

但在Vue3中，使用单独的Compistion Api也已经可以达到这样的逻辑复用的效果，并且定制性更强（手动控制那些需要loading的action）。因此在实际开发中，我并没有使用Pinia插件。

小结

本文首先展示了使用全局的reactive对象实现Vue全局状态的功能。

接着分析了Pinia的核心源码

createPinia返回pinia对象，单个pinia对象负责保存store列表
defineStore用来定义单个的Store，区分参数创建store对象，store上面包含一些内置方法如$pathch，也包含setup返回的state、getter和action
defineStore同时负责进行类型推断，将setup的返回值按照State、Getter、Action进行区分
storeToRefs只返回State和Getter

然后分析了pinia实现类型推断的相关代码。

最后实现了一个简易版本的mini-pinia，虽然代码很少，但感觉已经足够展示Pinia的核心原理。