From c8f727a0a2187c591134bd467efe426fb747ea40 Mon Sep 17 00:00:00 2001
From: Mikita Kamkou <44264841+nkeyy0@users.noreply.github.com>
Date: Tue, 2 Apr 2024 14:33:21 +0200
Subject: [PATCH] fix(types): fix storeToRefs state return type (#2574) (#2604)

* fix(types): Added new mapped type to keep original ref type in the storeToRefs return type (close #2574)

* test: add types test with custom ref type

* test: add type test for an options store that returns custom ref within the state function

* fix: Fix case when store has actions

* test: add type tests for stores with actions
---
 packages/pinia/src/storeToRefs.ts             |  33 ++++-
 .../pinia/test-dts/customizations.test-d.ts   | 117 ++++++++++++++++--
 2 files changed, 138 insertions(+), 12 deletions(-)

diff --git a/packages/pinia/src/storeToRefs.ts b/packages/pinia/src/storeToRefs.ts
index 76e91c333b..e00ccc15c7 100644
--- a/packages/pinia/src/storeToRefs.ts
+++ b/packages/pinia/src/storeToRefs.ts
@@ -11,7 +11,15 @@ import {
   toRefs,
 } from 'vue-demi'
 import { StoreGetters, StoreState } from './store'
-import type { PiniaCustomStateProperties, StoreGeneric } from './types'
+import type {
+  _ActionsTree,
+  _GettersTree,
+  _UnwrapAll,
+  PiniaCustomStateProperties,
+  StateTree,
+  Store,
+  StoreGeneric,
+} from './types'
 
 type ToComputedRefs<T> = {
   [K in keyof T]: ToRef<T[K]> extends Ref<infer U>
@@ -19,13 +27,30 @@ type ToComputedRefs<T> = {
     : ToRef<T[K]>
 }
 
+/**
+ * Extracts the refs of a state object from a store. If the state value is a Ref or type that extends ref, it will be kept as is.
+ * Otherwise, it will be converted into a Ref.
+ */
+declare type ToStateRefs<SS> =
+  SS extends Store<
+    string,
+    infer UnwrappedState,
+    _GettersTree<StateTree>,
+    _ActionsTree
+  >
+    ? UnwrappedState extends _UnwrapAll<Pick<infer State, infer Key>>
+      ? {
+          [K in Key]: ToRef<State[K]>
+        }
+      : ToRefs<UnwrappedState>
+    : ToRefs<StoreState<SS>>
+
 /**
  * Extracts the return type for `storeToRefs`.
  * Will convert any `getters` into `ComputedRef`.
  */
-export type StoreToRefs<SS extends StoreGeneric> = ToRefs<
-  StoreState<SS> & PiniaCustomStateProperties<StoreState<SS>>
-> &
+export type StoreToRefs<SS extends StoreGeneric> = ToStateRefs<SS> &
+  ToRefs<PiniaCustomStateProperties<StoreState<SS>>> &
   ToComputedRefs<StoreGetters<SS>>
 
 /**
diff --git a/packages/pinia/test-dts/customizations.test-d.ts b/packages/pinia/test-dts/customizations.test-d.ts
index ebf3b6435e..640ae466ad 100644
--- a/packages/pinia/test-dts/customizations.test-d.ts
+++ b/packages/pinia/test-dts/customizations.test-d.ts
@@ -121,13 +121,16 @@ expectType<{
 pinia.use(({ options, store }) => {
   const { debounce: debounceOptions } = options
   if (debounceOptions) {
-    return Object.keys(debounceOptions).reduce((debouncedActions, action) => {
-      debouncedActions[action] = debounce(
-        store[action],
-        debounceOptions[action]
-      )
-      return debouncedActions
-    }, {} as Record<string, (...args: any[]) => any>)
+    return Object.keys(debounceOptions).reduce(
+      (debouncedActions, action) => {
+        debouncedActions[action] = debounce(
+          store[action],
+          debounceOptions[action]
+        )
+        return debouncedActions
+      },
+      {} as Record<string, (...args: any[]) => any>
+    )
   }
 })
 
@@ -185,8 +188,106 @@ expectType<{
       const double = computed(() => n.value * 2)
       return {
         n,
-        double
+        double,
+      }
+    })()
+  )
+)
+
+expectType<{
+  n: Ref<number>
+  customN: Ref<number> & { plusOne: () => void }
+  double: ComputedRef<number>
+  myState: Ref<number>
+  stateOnly: Ref<number>
+}>(
+  storeToRefs(
+    defineStore('a', () => {
+      const n = ref(1)
+      const customN = ref(1) as Ref<number> & { plusOne: () => void }
+      const double = computed(() => n.value * 2)
+      return {
+        n,
+        customN,
+        double,
+      }
+    })()
+  )
+)
+
+expectType<{
+  n: Ref<number>
+  customN: Ref<number> & { plusOne: () => void }
+  double: ComputedRef<number>
+  myState: Ref<number>
+  stateOnly: Ref<number>
+}>(
+  storeToRefs(
+    defineStore('a', () => {
+      const n = ref(1)
+      const customN = ref(1) as Ref<number> & { plusOne: () => void }
+      const double = computed(() => n.value * 2)
+
+      function plusOne() {
+        customN.value++
+      }
+
+      return {
+        n,
+        customN,
+        double,
+        plusOne,
       }
     })()
   )
 )
+
+expectType<{
+  n: Ref<number>
+  customN: Ref<number> & { plusOne: () => void }
+  double: ComputedRef<number>
+  myState: Ref<number>
+  stateOnly: Ref<number>
+}>(
+  storeToRefs(
+    defineStore('a', {
+      state: () => ({
+        n: 1,
+        customN: ref(1) as Ref<number> & { plusOne: () => void },
+      }),
+      getters: {
+        double: (state) => state.n * 2,
+      },
+      actions: {
+        plusOne() {
+          this.n++
+        },
+      },
+    })()
+  )
+)
+
+expectType<{
+  n: Ref<number>
+  customN: Ref<number> & { plusOne: () => void }
+  double: ComputedRef<number>
+  myState: Ref<number>
+  stateOnly: Ref<number>
+}>(
+  storeToRefs(
+    defineStore('a', {
+      state: () => ({
+        n: 1,
+        customN: ref(1) as Ref<number> & { plusOne: () => void },
+      }),
+      getters: {
+        double: (state) => state.n * 2,
+      },
+      actions: {
+        plusOne() {
+          this.n++
+        },
+      },
+    })()
+  )
+)