From 34cef9b9b223a9132e5a1745bc858a19d91af130 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                      | 22 +++++++++++++++++-----
 test-dts/customizations.test-d.ts | 12 ++++++++++++
 4 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index a8da3e30..f5a9b2c6 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -17,6 +17,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 b529e2ff..ae2d91cd 100644
--- a/src/rootStore.ts
+++ b/src/rootStore.ts
@@ -8,6 +8,7 @@ import {
   Store,
   DefineStoreOptions,
   ActionsTree,
+  PiniaCustomStateProperties,
 } from './types'
 import type Vue from 'vue'
 
@@ -68,7 +69,9 @@ 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 6ba83e9a..b17239c8 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -198,7 +198,7 @@ export interface StoreWithState<Id extends string, S extends StateTree> {
   /**
    * 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.
@@ -339,7 +339,8 @@ export type Store<
   UnwrapRef<S> &
   StoreWithGetters<G> &
   StoreWithActions<A> &
-  PiniaCustomProperties<Id, S, G, A>
+  PiniaCustomProperties<Id, S, G, A> &
+  PiniaCustomStateProperties<S>
 
 // TODO: check if it's possible to add = to StoreDefinition and Store and cleanup GenericStore and the other one
 
@@ -381,6 +382,11 @@ export interface PiniaCustomProperties<
   A = 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
  *
@@ -388,7 +394,7 @@ export interface PiniaCustomProperties<
  */
 export type GettersTree<S extends StateTree> = Record<
   string,
-  ((state: UnwrapRef<S>) => any) | (() => any)
+  ((state: UnwrapRef<S & PiniaCustomStateProperties<S>>) => any) | (() => any)
 >
 
 /**
@@ -418,7 +424,12 @@ export interface DefineStoreOptions<
    * Optional object of getters.
    */
   getters?: G &
-    ThisType<UnwrapRef<S> & StoreWithGetters<G> & PiniaCustomProperties>
+    ThisType<
+      UnwrapRef<S> &
+        StoreWithGetters<G> &
+        PiniaCustomProperties &
+        PiniaCustomStateProperties
+    >
   /**
    * Optional object of actions.
    */
@@ -428,6 +439,7 @@ export interface DefineStoreOptions<
         UnwrapRef<S> &
         StoreWithState<Id, S> &
         StoreWithGetters<G> &
-        PiniaCustomProperties
+        PiniaCustomProperties &
+        PiniaCustomStateProperties
     >
 }
diff --git a/test-dts/customizations.test-d.ts b/test-dts/customizations.test-d.ts
index b9e69874..30afff13 100644
--- a/test-dts/customizations.test-d.ts
+++ b/test-dts/customizations.test-d.ts
@@ -9,6 +9,10 @@ declare module '../dist/src' {
     $actions: Array<keyof A>
   }
 
+  export interface PiniaCustomStateProperties<S> {
+    myState: number
+  }
+
   export interface DefineStoreOptions<Id, S, G, A> {
     debounce?: {
       // Record<keyof A, number>
@@ -23,6 +27,9 @@ pinia.use((context) => {
   expectType<string>(context.options.id)
   expectType<string>(context.store.$id)
 
+  expectType<number>(context.store.$state.myState)
+  expectType<number>(context.store.myState)
+
   return {
     $actions: Object.keys(context.options.actions || {}),
   }
@@ -30,10 +37,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()
@@ -47,6 +57,8 @@ const useStore = defineStore({
   },
 })
 
+expectType<{ myState: number }>(useStore().$state)
+
 type Procedure = (...args: any[]) => any
 
 function debounce<F extends Procedure>(fn: F, time = 200) {
-- 
2.47.3