fix(types): unwrap refs in mapWritableState for setup stores

Author: jh-leong

packages/pinia/src/mapHelpers.ts

@@ -465,10 +465,11 @@ export function mapWritableState<
   G extends _GettersTree<S>,
   A,
   KeyMapper extends Record<string, keyof S>,
+  SS extends StateTree = Store<Id, S, G, A>,
 >(
   useStore: StoreDefinition<Id, S, G, A>,
   keyMapper: KeyMapper
-): _MapWritableStateObjectReturn<S, KeyMapper>
+): Pick<_MapWritableStateObjectReturn<SS, KeyMapper>, keyof KeyMapper>
 /**
  * Allows using state and getters from one store without using the composition
  * API (`setup()`) by generating an object to be spread in the `computed` field
@@ -483,15 +484,11 @@ export function mapWritableState<
   G extends _GettersTree<S>,
   A,
   Keys extends keyof S,
+  SS extends StateTree = Store<Id, S, G, A>,
 >(
   useStore: StoreDefinition<Id, S, G, A>,
   keys: readonly Keys[]
-): {
-  [K in Keys]: {
-    get: () => S[K]
-    set: (value: S[K]) => any
-  }
-}
+): Pick<_MapWritableStateReturn<SS>, Keys>
 /**
  * Allows using state and getters from one store without using the composition
  * API (`setup()`) by generating an object to be spread in the `computed` field

packages/pinia/src/types.ts

@@ -11,7 +11,7 @@ import { Pinia } from './rootStore'
 /**
  * Generic state of a Store
  */
-export type StateTree = Record<string | number | symbol, any>
+export type StateTree = Record<PropertyKey, any>
 
 export function isPlainObject<S extends StateTree>(
   value: S | unknown

packages/pinia/test-dts/mapHelpers.test-d.ts

@@ -116,7 +116,7 @@ expectType<{
   }
 }>(mapWritableState(useOptionsStore, ['a']))
 // @ts-expect-error: only defined in array
-mapWritableState(useStore, ['a']).b
+mapWritableState(useOptionsStore, ['a']).b
 
 expectType<{
   newA: {
@@ -126,9 +126,23 @@ expectType<{
 }>(mapWritableState(useOptionsStore, { newA: 'a' }))
 
 // @ts-expect-error: cannot use a getter
-mapWritableState(useStore, ['upper'])
+mapWritableState(useOptionsStore, ['upper'])
 // @ts-expect-error: cannot use a getter
-mapWritableState(useStore, { up: 'upper' })
+mapWritableState(useOptionsStore, { up: 'upper' })
+
+expectType<{
+  foo: {
+    get: () => 'on' | 'off'
+    set: (v: 'on' | 'off') => any
+  }
+}>(mapWritableState(useSetupStore, { foo: 'a' }))
+
+expectType<{
+  a: {
+    get: () => 'on' | 'off'
+    set: (v: 'on' | 'off') => any
+  }
+}>(mapWritableState(useSetupStore, ['a']))
 
 const setupStoreWithState = mapState(useSetupStore, ['a'])