From 391f9ac4f5e845afdd4b7d44fb1a9c7fb2b299a0 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Fri, 13 Oct 2023 11:22:53 +0200
Subject: [PATCH] feat(types): SetupStoreDefinition

---
 packages/pinia/src/index.ts                   |  7 +++-
 packages/pinia/src/store.ts                   | 14 ++++++++
 packages/pinia/test-dts/typeHelpers.test-d.ts | 33 ++++++++++++++-----
 3 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/packages/pinia/src/index.ts b/packages/pinia/src/index.ts
index 9735ef3767..fc28044f4d 100644
--- a/packages/pinia/src/index.ts
+++ b/packages/pinia/src/index.ts
@@ -12,7 +12,12 @@ export type {
 } from './rootStore'
 
 export { defineStore, skipHydrate } from './store'
-export type { StoreActions, StoreGetters, StoreState } from './store'
+export type {
+  StoreActions,
+  StoreGetters,
+  StoreState,
+  SetupStoreDefinition,
+} from './store'
 
 export type {
   StateTree,
diff --git a/packages/pinia/src/store.ts b/packages/pinia/src/store.ts
index e76381f041..cf1d5be173 100644
--- a/packages/pinia/src/store.ts
+++ b/packages/pinia/src/store.ts
@@ -964,3 +964,17 @@ export function defineStore(
 
   return useStore
 }
+
+/**
+ * Return type of `defineStore()` with a setup function.
+ * - `Id` is a string literal of the store's name
+ * - `SS` is the return type of the setup function
+ * @see {@link StoreDefinition}
+ */
+export interface SetupStoreDefinition<Id extends string, SS>
+  extends StoreDefinition<
+    Id,
+    _ExtractStateFromSetupStore<SS>,
+    _ExtractGettersFromSetupStore<SS>,
+    _ExtractActionsFromSetupStore<SS>
+  > {}
diff --git a/packages/pinia/test-dts/typeHelpers.test-d.ts b/packages/pinia/test-dts/typeHelpers.test-d.ts
index 2bd7c9eb37..5e1a432d42 100644
--- a/packages/pinia/test-dts/typeHelpers.test-d.ts
+++ b/packages/pinia/test-dts/typeHelpers.test-d.ts
@@ -1,6 +1,7 @@
-import { StoreDefinition } from './'
-import { computed, ref } from 'vue'
+import { ComputedRef, Ref, computed, ref } from 'vue'
 import {
+  StoreDefinition,
+  SetupStoreDefinition,
   StoreState,
   StoreGetters,
   StoreActions,
@@ -60,10 +61,26 @@ expectType<{ n: number }>(storeState(useOptionsStore))
 
 expectType<{ double: number }>(storeGetters(useOptionsStore))
 
-expectType<{ n: number }>(
-  storeState(
-    defineStore('', {
-      state: () => ({ n: ref(0) }),
-    })
-  )
+expectType<
+  SetupStoreDefinition<
+    'a',
+    {
+      n: Ref<number>
+      double: ComputedRef<number>
+      increment: () => void
+    }
+  >
+>(
+  defineStore('a', () => {
+    const n = ref(0)
+    const double = computed(() => n.value * 2)
+    function increment() {
+      n.value++
+    }
+    return {
+      double,
+      increment,
+      n,
+    }
+  })
 )