feat(patch): allow passing a function

vuejs · Apr 1, 2021 · 8d545e4 · 8d545e4
1 parent 0265efa
commit 8d545e4
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 3 deletions.
diff --git a/__tests__/store.patch.spec.ts b/__tests__/store.patch.spec.ts
@@ -12,6 +12,7 @@ describe('store.$patch', () => {
           foo: 'foo',
           a: { b: 'string' },
         },
+        list: [] as number[],
       }),
     })()
   }
@@ -25,6 +26,23 @@ describe('store.$patch', () => {
         foo: 'foo',
         a: { b: 'string' },
       },
+      list: [],
+    })
+  })
+
+  it('patches using a function', () => {
+    const store = useStore()
+    store.$patch((state) => {
+      state.a = !state.a
+      state.list.push(1)
+    })
+    expect(store.$state).toEqual({
+      a: false,
+      nested: {
+        foo: 'foo',
+        a: { b: 'string' },
+      },
+      list: [1],
     })
   })
 
@@ -37,6 +55,7 @@ describe('store.$patch', () => {
         foo: 'bar',
         a: { b: 'string' },
       },
+      list: [],
     })
     store.$patch({ nested: { a: { b: 'hello' } } })
     expect(store.$state).toEqual({
@@ -45,6 +64,7 @@ describe('store.$patch', () => {
         foo: 'bar',
         a: { b: 'hello' },
       },
+      list: [],
     })
   })
 
@@ -57,6 +77,7 @@ describe('store.$patch', () => {
         foo: 'hello',
         a: { b: 'string' },
       },
+      list: [],
     })
   })
 })
diff --git a/src/store.ts b/src/store.ts
@@ -89,14 +89,27 @@ function initStore<Id extends string, S extends StateTree>(
   let isListening = true
   let subscriptions: SubscriptionCallback<S>[] = []
 
-  function $patch(partialState: DeepPartial<S>): void {
+  function $patch(stateMutation: (state: S) => void): void
+  function $patch(partialState: DeepPartial<S>): void
+  function $patch(
+    partialStateOrMutator: DeepPartial<S> | ((state: S) => void)
+  ): void {
+    let partialState: DeepPartial<S> = {}
+    let type: string
     isListening = false
-    innerPatch(pinia.state.value[$id], partialState)
+    if (typeof partialStateOrMutator === 'function') {
+      partialStateOrMutator(pinia.state.value[$id])
+      type = '🧩 patch'
+    } else {
+      innerPatch(pinia.state.value[$id], partialStateOrMutator)
+      partialState = partialStateOrMutator
+      type = '⤵️ patch'
+    }
     isListening = true
     // because we paused the watcher, we need to manually call the subscriptions
     subscriptions.forEach((callback) => {
       callback(
-        { storeName: $id, type: '⤵️ patch', payload: partialState },
+        { storeName: $id, type, payload: partialState },
         pinia.state.value[$id]
       )
     })

diff --git a/src/types.ts b/src/types.ts
@@ -71,6 +71,15 @@ export interface StoreWithState<Id extends string, S extends StateTree> {
    */
   $patch(partialState: DeepPartial<S>): void
 
+  /**
+   * Group multiple changes into one function. Useful when mutating objects like
+   * Sets or arrays and applying an object patch isn't practical, e.g. appending
+   * to an array.
+   *
+   * @param stateMutator - function that mutates `state`
+   */
+  $patch(stateMutator: (state: S) => void): void
+
   /**
    * Resets the store to its initial state by removing all subscriptions and
    * building a new state object