From 17fcbcafd30eec3af609bcc4549eee08968ea1a2 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Sat, 15 May 2021 17:15:02 +0200
Subject: [PATCH] feat(types): allow defining custom state properties

---
 src/index.ts                      |  1 +
 src/rootStore.ts                  |  5 ++++-
 src/types.ts                      | 12 +++++++++---
 test-dts/customizations.test-d.ts | 10 ++++++++++
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 27892b70..7c577382 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -18,6 +18,7 @@ export type {
   StoreOnActionListenerContext,
   SubscriptionCallback,
   PiniaCustomProperties,
+  PiniaCustomStateProperties,
   DefineStoreOptions,
 } from './types'
 export { MutationType } from './types'
diff --git a/src/rootStore.ts b/src/rootStore.ts
index e991852b..37dc3dd4 100644
--- a/src/rootStore.ts
+++ b/src/rootStore.ts
@@ -9,6 +9,7 @@ import {
   Store,
   GettersTree,
   ActionsTree,
+  PiniaCustomStateProperties,
 } from './types'
 
 /**
@@ -153,5 +154,7 @@ export interface PiniaStorePlugin {
    *
    * @param context - Context
    */
-  (context: PiniaPluginContext): Partial<PiniaCustomProperties> | void
+  (context: PiniaPluginContext): Partial<
+    PiniaCustomProperties & PiniaCustomStateProperties
+  > | void
 }
diff --git a/src/types.ts b/src/types.ts
index 66c4d6f0..b83fa2c3 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -231,7 +231,7 @@ export interface StoreWithState<
   /**
    * State of the Store. Setting it will replace the whole state.
    */
-  $state: UnwrapRef<S>
+  $state: UnwrapRef<S> & PiniaCustomStateProperties<S>
 
   /**
    * Private property defining the pinia the store is attached to.
@@ -380,7 +380,8 @@ export type Store<
   UnwrapRef<S> &
   StoreWithGetters<G> &
   StoreWithActions<A> &
-  PiniaCustomProperties<Id, S, G, A>
+  PiniaCustomProperties<Id, S, G, A> &
+  PiniaCustomStateProperties<S>
 
 /**
  * Return type of `defineStore()`. Function that allows instantiating a store.
@@ -420,6 +421,11 @@ export interface PiniaCustomProperties<
   A /* extends ActionsTree */ = ActionsTree
 > {}
 
+/**
+ * Properties that are added to every `store.$state` by `pinia.use()`
+ */
+export interface PiniaCustomStateProperties<S extends StateTree = StateTree> {}
+
 /**
  * Type of an object of Getters that infers the argument
  *
@@ -427,7 +433,7 @@ export interface PiniaCustomProperties<
  */
 export type GettersTree<S extends StateTree> = Record<
   string,
-  ((state: UnwrapRef<S>) => any) | (() => any)
+  ((state: UnwrapRef<S & PiniaCustomStateProperties<S>>) => any) | (() => any)
 >
 
 /**
diff --git a/test-dts/customizations.test-d.ts b/test-dts/customizations.test-d.ts
index c4aaf6d4..e37c4972 100644
--- a/test-dts/customizations.test-d.ts
+++ b/test-dts/customizations.test-d.ts
@@ -11,6 +11,10 @@ declare module '../dist/pinia' {
     $actions: Array<keyof A>
   }
 
+  export interface PiniaCustomStateProperties<S> {
+    myState: number
+  }
+
   export interface DefineStoreOptions<Id, S, G, A> {
     debounce?: {
       // Record<keyof A, number>
@@ -26,6 +30,9 @@ pinia.use((context) => {
   expectType<string>(context.store.$id)
   expectType<App>(context.app)
 
+  expectType<number>(context.store.$state.myState)
+  expectType<number>(context.store.myState)
+
   return {
     $actions: Object.keys(context.options.actions || {}),
   }
@@ -33,10 +40,13 @@ pinia.use((context) => {
 
 const useStore = defineStore({
   id: 'main',
+  state: () => ({}),
   actions: {
     one() {},
     two() {
       this.one()
+      expectType<number>(this.$state.myState)
+      expectType<number>(this.myState)
     },
     three() {
       this.two()
-- 
2.47.2