手撕pinia源码（二）：实现defineStore

手撕pinia源码目录

手撕pinia源码（一）： pinia的使用

手撕pinia源码（二）：实现defineStore

手撕pinia源码（三）：实现pinia内置方法

源码地址传送门

前期准备

安装vue

根据vue官网快速搭建一个项目，删掉不必要的文件，以免影响我们开发，只保留最简洁的功能页面即可。

创建pinia

在/src目录下创建pinia目录，并新建一个入口文件index.js。新建createPinia.js文件，此文件提供createPinia方法是插件在vue中注册的方法。新建store.js文件，此文件提供defineStore方法，用于定义一个store仓库的方法。

// createPinia.js
export function createPinia() {
}

// store.js
export function defineStore() {
}

// index.js
import { createPinia } from "./createPinia.js";
import { defineStore } from './store';

export {
    createPinia,
    defineStore
}

这就是pinia的基本结构，index.js直接导出createPinia、defineStore方法供外部使用即可。

接下来修改main.js注册pinia插件，报错没关系，因为我们还没有实现createPinia方法。

// main.js
import './assets/main.css'

import { createApp } from 'vue'
import { createPinia } from '@/pinia'
import App from './App.vue'

const app = createApp(App)

app.use(createPinia())

app.mount('#app')

最终的项目目录结构如下：

.
├── README.md
├── index.html
├── jsconfig.json
├── package.json
├── src
│   ├── App.vue
│   ├── assets
│   │   ├── base.css
│   │   ├── logo.svg
│   │   └── main.css
│   ├── main.js
│   ├── pinia
│   │   ├── createPinia.js
│   │   ├── index.js
│   │   └── store.js
└── vite.config.js

实现createPinia方法

vue插件规定必须实现一个install方法，install方法能够获取到当前的app实例，因此我们修改一下createPinia.js

// createPinia.js
export function createPinia() {
    const pinia = {
        install(app) {
        }
    }
    return pinia;
}

源码分析

1. pinia中实例总是全局的，所以我们需要给app挂载全的pinia实例。
2. 源码中_s，pinia支持多仓库，而他的仓库存储的数据结构大致为{id1: store1, id2: store2}这种形式，因此我们需要创建一个对象存储所有的仓库。
3. 源码中_e是一个effectScope的vueAPI，官方解释effectScope，它可以收集所有的副作用函数通过调用stop来让所有的响应式数据停止响应，后续我们再详细介绍它是干什么用的，而副作用函数可谓是vue响应式的核心，可以理解为如果关闭了副作用函数那么数据变化，视图也不会变化。
4. 源码中state是一个ref的响应式对象，其主要用于存储所有store的状态值。

编写代码

// createPinia.js
import { effectScope, ref } from 'vue';
import { piniaSymbol } from './rootStore';

export function createPinia() {
    const scope = effectScope();
    const state = scope.run(() => ref({})); // 用于存储每个store的state方法
    const pinia = {
        _s: new Map(), // 用来存储所有的store实例，{storeName: store, storeName2: store2}
        _e: scope, // effectScope实例
        state, // 存储所有store的state状态值
        install(app) {
            // 注入依赖
            app.provide(piniaSymbol, pinia); // 让所有的store都能够拿到pinia实例

            app.config.globalProperties.$pinia = pinia;
        }
    }
    return pinia;
}

// rootStore
export const piniaSymbol = Symbol();

• 代码6行，创建一个effectScope，将该实例挂载到pinia实例中，命名为_e。
• 代码7行，创建一个ref响应式对象，将对象挂载到pinia实例中，命名为state。
• 代码9行，创建一个Map对象，store存储形式是{id1: store1, id2: store2}的，所以id不允许重复，因此用Map来做存储会方便许多。
• 代码14行，插件被注册时，需要在vue全局属性中挂载$pinia，同时provide注入依赖，提供给全局使用，名称跟官方一样用symbol值。

createPinia方法就基本完成了，是不是非常简单呢，接着进入到核心代码的编写defineStore。

实现defineStore方法

源码分析

1. defineStore有三个入参，分别是idOrOptions、setup、setupOptions，上一章我们讲到，defineStore的三种用法。
2. 可以从源码1678行看到，默认是有id和options选项，通过类型判断用户入参进行相关的处理。
3. 源码1747行，返回了useStore实例给用户调用，再看看useStore的实现，源码1697行注入依赖pinia实例，然后在源码1706行获取pinia实例的_s是否有保存id这个store，如果没有就进行初始化，初始化分情况调用了createSetupStore和createOptionsStore两种方法。

编写代码

// store.js
export function defineStore(idOrOptions, setup) {
    let id, options;
    const isSetupStore = typeof setup === 'function';
    // 如果第一个参数是string,那么它就是id
    if (typeof idOrOptions === 'string') {
        id = idOrOptions;
        options = setup;
    } else {
        id = idOrOptions.id;
        options = idOrOptions;
    }

    function useStore() {
        // 获取当前组件实例，拿到pinia实例
        const instance = getCurrentInstance();
        const pinia = instance && inject(piniaSymbol);
        // 判断是否初始化，如果pinia._s没有这个id，则设置一个
        if (!pinia._s.has(id)) {
            // 判断第二个参数是否是一个函数，如果是函数那么则是使用compositio形式
            if (isSetupStore) {
                createSetupStore(id, setup, pinia);
            } else {
                createOptionsStore(id, options, pinia);
            }
        }
        const store = pinia._s.get(id);
        return store;
    }
    return useStore;
}

• 代码4行，判断第二个参数是否为一个函数，如果是函数，则是compositionAPI,如果不是函数，则是optionsAPI。
• 代码6行，判断第一个参数是否为string类型，如果是，则它就是该store的唯一值id，否则从第二个参数中取id。
• 代码14行，实现一个useStore方法返回store实例。
• 代码17行，通过getCurrentInstanceAPI获取当前的vue实例，如果有实例，则注入依赖piniaSymbol获取到pinia实例。
• 代码19行，从pinia._s中通过id取store值，如果没取到，进行初始化。
• 代码21行，根据isSetupStore判断调用createSetupStore方法还是createOptionsStore方法。
• 代码27行，获取到store实例，用户调用useStore时就能获取到对应的store。

实现createSetupStore和createOptionsStore方法

createSetupStore源码分析

createSetupStore源码有点长，我们截取核心部分分析。

1. 源码1403，partialStore是pinia自带的方法和属性，我们在下一章会实现这里面的方法，store变量是将原属性和用户定义的属性和方法合并。
2. 源码1440，将store存储到pinia._s下，方便后面二次读取就无需再进行初始化。
3. 源码1443，将用户传入的setup也是用effectScope包一层。
4. 源码1445，循环用户传入的setup，因为用户传入的setup是散乱的，他和options不同，需要判断用户到底写的是state状态，还是方法，又或者它是一个computed属性，所以需要循环对这些进行处理
5. 循环内容大致逻辑就是，将所有的状态存储到pinia.state中，将所有函数重新绑定this，因为如果用户将方法结构出来使用的话，this就会错误或丢失，例如: const { increment } = useCounterStore()，这时调用increment时，this并不是指向store无法取到值。

createSetupStore编写代码

// store.js
import { getCurrentInstance, inject, reactive, effectScope, computed, isRef, isReactive } from 'vue';
function isComputed(v) {
    return (isRef(v) && v.effect);
}
function createSetupStore(id, setup, pinia, isOptions) {
    let scope;
    const partialStore = reactive({});
    const initState = pinia.state.value[id];
    if (!initState) {
        pinia.state.value[id] = {};
    }
    const setupStore = pinia._e.run(() => {
        scope = effectScope();
        return scope.run(() => setup());
    });

    function warpAction(name, action) {
        return function () {
            let res = action.apply(partialStore, arguments);
            return res;
        }
    }
    for(let key in setupStore) {
        const prop = setupStore[key];
        if (typeof prop === 'function') {
            setupStore[key] = warpAction(key, prop);
        }
        if (!isOptions) {
            // 如果setup API 需要拿到状态存到全局的state中, computed也是ref，需要另外处理
            if ((isRef(prop) && !isComputed(prop)) || isReactive(prop)) {
                pinia.state.value[id][key] = prop;
            }
        }
    }
    pinia._s.set(id, partialStore);
    Object.assign(partialStore, setupStore);
    return partialStore;
}

• 代码8行，reactive实现一个响应式对象，用于存储pinia的内置方法。
• 代码9行，先通过pinia.state获取state状态，如果没有，进行初始化。
• 代码11行，给定一个pinia.state的默认值为{}。
• 代码13行，调用pinia._e执行run方法，然后再定一个当前store的effectScope方法，执行用户传入的setup函数，effectScope是可以嵌套使用的，可以假设为pinia._e是全局的控制的是整个pinia的响应式，而当前的effectScope是当前store局部的，控制的是当前store的响应式。
• 代码24行，循环所有用户定义的setup，判断如果是function的话，调用warpAction重新绑定this的指向。
• 代码32行，需要区分哪些是用户的状态，我们要先了解，const a = ref(a);这种是用户定义的状态，所以我们可以调用isRef、isReactive来判断。但是有一点需要注意，const b = computed()，此时isRef(b) === true，就是说computed也是ref，我们要另外判断是不是computed属性。
• 代码3行，判断一个属性是否为computed，首先它一定是ref，其次，判断computed下是否有effect这个方法，这里关于vue源码不多赘述。
• 代码36行，将用户定义的store和pinia内置属性方法合并，存储到pinia._s中，后面二次调用useStore时无需再初始化，直接取值就能返回。

createOptionsStore源码分析

1. 源码1196行，初始化pinia.state的属性，源码1204行，执行用户传入的state方法（optionsAPI中，所有的state是函数，getter是对象，action是对象）将其存储到全局的state。
2. 源码1209行，获取用户定义的state属性，要知道用户定义只是一个普通的值，并不具有响应式，所以需要toRefs让所有对象转成响应式。
3. 源码1217行，处理用户定义的getter，用户定义时也是一个普通函数，所以也需要将其处理为computed，源码1228行.call绑定的this一定要指向自己sotre（通过pinia._s获取，不用担心获取不到，因为computed是用户取值时才执行，所以pinia._s已经存在。
4. 代码1233行，调用刚刚的createSetupStore方法，可以看到其实options和composition都用同一套逻辑，只是当用户使用options时，我们给他重新组装一个setup然后交给createSetupStore函数处理。

createOptionsStore编写代码

// store.js
// 创建选项式的store
function createOptionsStore(id, options, pinia) {
    const { state, getters, actions } = options;

    function setup() {
        const localState = pinia.state.value[id] = state ? state() : {};
        return Object.assign(toRefs(ref(localState).value), actions, Object.keys(getters).reduce((memo, name) => {
            memo[name] = computed(() => {
                let store = pinia._s.get(id);
                return getters[name].call(store, store);
            })
            return memo;
        }, {}));
    }
    return createSetupStore(id, setup, pinia, true); 
}

• 代码7行，将所有的状态存储到全局的pinia.state中。
• 代码8行，处理用户的state和getter，state转成ref，而getter则循环将所有的属性重新用computed进行绑定即可，跟源码实现差不多。

测试demo

// App.vue
<template>
    <div>
        <div>{{ store.counter }} / {{ store.dobuleCount }}</div>
        <button @click="handleClick">optionsStore</button>
    </div>
    <div>
        <div>{{ store2.counter }}/ {{ store2.dobuleCount }}</div>
        <button @click="handleClick2">setupStore</button>
    </div>
</template>

<script setup>
import { useCounterStore } from './stores/counter';
import { useCounterStore2 } from './stores/counter2';

const store = useCounterStore();
const store2 = useCounterStore2();

function handleClick() {
    store.increment();
}
function handleClick2() {
    store2.increment();
}
</script>

options形式定义store

// store/counter.js
import { defineStore } from '@/pinia'

export const useCounterStore = defineStore('counterStore', {
  state: () => {
    return {
      counter: 0
    }
  },
  actions: {
    increment() {
      this.counter++
    }    
  },
  getters: {
    dobuleCount() {
      console.log(this.counter, 111)
      return this.counter * 2
    }
  }
})

compositon形式定义store

import { ref, computed } from 'vue'
import { defineStore } from '@/pinia'

export const useCounterStore2 = defineStore('counterStore2', () => {
  const counter = ref(0);
  function increment() {
    this.counter++
  }
  const dobuleCount = computed(() => {
    return counter.value * 2;
  })
  return { counter, increment, dobuleCount }
})

npm run dev运行代码看看效果。

结束

我们已经把pinia的核心功能都实现了一遍，如果有不懂的，欢迎留言或移步至文章开头查看完整源码，跟着逻辑打印一遍就会觉得pinia非常简单，下一章将讲解pinia内置的方法的实现。