add state extensions from enhancers

Author: ben.durrant

packages/toolkit/src/configureStore.ts

@@ -24,6 +24,7 @@ import type {
   NoInfer,
   ExtractDispatchExtensions,
   ExtractStoreExtensions,
+  ExtractStateExtensions,
 } from './tsHelpers'
 import { EnhancerArray } from './utils'
 
@@ -129,7 +130,8 @@ export type EnhancedStore<
   A extends Action = AnyAction,
   M extends Middlewares<S> = Middlewares<S>,
   E extends Enhancers = Enhancers
-> = ToolkitStore<S, A, M> & ExtractStoreExtensions<E>
+> = ToolkitStore<S & ExtractStateExtensions<E>, A, M> &
+  ExtractStoreExtensions<E>
 
 /**
  * A friendly abstraction over the standard Redux `createStore()` function.

packages/toolkit/src/tests/EnhancerArray.typetest.ts

@@ -3,13 +3,19 @@ import type { StoreEnhancer } from 'redux'
 
 declare const expectType: <T>(t: T) => T
 
-declare const enhancer1: StoreEnhancer<{
-  has1: true
-}>
+declare const enhancer1: StoreEnhancer<
+  {
+    has1: true
+  },
+  { stateHas1: true }
+>
 
-declare const enhancer2: StoreEnhancer<{
-  has2: true
-}>
+declare const enhancer2: StoreEnhancer<
+  {
+    has2: true
+  },
+  { stateHas2: true }
+>
 
 {
   // prepend single element
@@ -19,9 +25,12 @@ declare const enhancer2: StoreEnhancer<{
       enhancers: (dE) => dE.prepend(enhancer1),
     })
     expectType<true>(store.has1)
+    expectType<true>(store.getState().stateHas1)
 
     // @ts-expect-error
     expectType<true>(store.has2)
+    // @ts-expect-error
+    expectType<true>(store.getState().stateHas2)
   }
 
   // prepend multiple (rest)
@@ -31,10 +40,14 @@ declare const enhancer2: StoreEnhancer<{
       enhancers: (dE) => dE.prepend(enhancer1, enhancer2),
     })
     expectType<true>(store.has1)
+    expectType<true>(store.getState().stateHas1)
     expectType<true>(store.has2)
+    expectType<true>(store.getState().stateHas2)
 
     // @ts-expect-error
     expectType<true>(store.has3)
+    // @ts-expect-error
+    expectType<true>(store.getState().stateHas3)
   }
 
   // prepend multiple (array notation)
@@ -44,10 +57,14 @@ declare const enhancer2: StoreEnhancer<{
       enhancers: (dE) => dE.prepend([enhancer1, enhancer2] as const),
     })
     expectType<true>(store.has1)
+    expectType<true>(store.getState().stateHas1)
     expectType<true>(store.has2)
+    expectType<true>(store.getState().stateHas2)
 
     // @ts-expect-error
     expectType<true>(store.has3)
+    // @ts-expect-error
+    expectType<true>(store.getState().stateHas3)
   }
 
   // concat single element
@@ -57,9 +74,12 @@ declare const enhancer2: StoreEnhancer<{
       enhancers: (dE) => dE.concat(enhancer1),
     })
     expectType<true>(store.has1)
+    expectType<true>(store.getState().stateHas1)
 
     // @ts-expect-error
     expectType<true>(store.has2)
+    // @ts-expect-error
+    expectType<true>(store.getState().stateHas2)
   }
 
   // prepend multiple (rest)
@@ -69,10 +89,14 @@ declare const enhancer2: StoreEnhancer<{
       enhancers: (dE) => dE.concat(enhancer1, enhancer2),
     })
     expectType<true>(store.has1)
+    expectType<true>(store.getState().stateHas1)
     expectType<true>(store.has2)
+    expectType<true>(store.getState().stateHas2)
 
     // @ts-expect-error
     expectType<true>(store.has3)
+    // @ts-expect-error
+    expectType<true>(store.getState().stateHas3)
   }
 
   // concat multiple (array notation)
@@ -82,10 +106,14 @@ declare const enhancer2: StoreEnhancer<{
       enhancers: (dE) => dE.concat([enhancer1, enhancer2] as const),
     })
     expectType<true>(store.has1)
+    expectType<true>(store.getState().stateHas1)
     expectType<true>(store.has2)
+    expectType<true>(store.getState().stateHas2)
 
     // @ts-expect-error
     expectType<true>(store.has3)
+    // @ts-expect-error
+    expectType<true>(store.getState().stateHas3)
   }
 
   // concat and prepend
@@ -95,9 +123,13 @@ declare const enhancer2: StoreEnhancer<{
       enhancers: (dE) => dE.concat(enhancer1).prepend(enhancer2),
     })
     expectType<true>(store.has1)
+    expectType<true>(store.getState().stateHas1)
     expectType<true>(store.has2)
+    expectType<true>(store.getState().stateHas2)
 
     // @ts-expect-error
     expectType<true>(store.has3)
+    // @ts-expect-error
+    expectType<true>(store.getState().stateHas3)
   }
 }

packages/toolkit/src/tests/configureStore.typetest.ts

@@ -209,6 +209,72 @@ const _anyMiddleware: any = () => () => () => {}
     expectType<string>(storeWithCallback.someProperty)
     expectType<number>(storeWithCallback.anotherProperty)
   }
+
+  {
+    type StateExtendingEnhancer = StoreEnhancer<{}, { someProperty: string }>
+
+    const someStateExtendingEnhancer: StateExtendingEnhancer =
+      (next) =>
+      // @ts-expect-error how do you properly return an enhancer that extends state?
+      (...args) => {
+        const store = next(...args)
+        const getState = () => ({
+          ...store.getState(),
+          someProperty: 'some value',
+        })
+        return {
+          ...store,
+          getState,
+        }
+      }
+
+    type AnotherStateExtendingEnhancer = StoreEnhancer<
+      {},
+      { anotherProperty: number }
+    >
+
+    const anotherStateExtendingEnhancer: AnotherStateExtendingEnhancer =
+      (next) =>
+      // @ts-expect-error any input on this would be great
+      (...args) => {
+        const store = next(...args)
+        const getState = () => ({
+          ...store.getState(),
+          anotherProperty: 123,
+        })
+        return {
+          ...store,
+          getState,
+        }
+      }
+
+    const store = configureStore({
+      reducer: () => ({ aProperty: 0 }),
+      enhancers: [
+        someStateExtendingEnhancer,
+        anotherStateExtendingEnhancer,
+        // this doesn't work without the as const
+      ] as const,
+    })
+
+    const state = store.getState()
+
+    expectType<number>(state.aProperty)
+    expectType<string>(state.someProperty)
+    expectType<number>(state.anotherProperty)
+
+    const storeWithCallback = configureStore({
+      reducer: () => ({ aProperty: 0 }),
+      enhancers: (dE) =>
+        dE.concat(someStateExtendingEnhancer, anotherStateExtendingEnhancer),
+    })
+
+    const stateWithCallback = storeWithCallback.getState()
+
+    expectType<number>(stateWithCallback.aProperty)
+    expectType<string>(stateWithCallback.someProperty)
+    expectType<number>(stateWithCallback.anotherProperty)
+  }
 }
 
 /**

packages/toolkit/src/tsHelpers.ts

@@ -119,7 +119,34 @@ export type ExtractStoreExtensions<E> = E extends EnhancerArray<
   ? UnionToIntersection<
       E[number] extends StoreEnhancer<infer Ext>
         ? Ext extends {}
-          ? Ext
+          ? IsAny<Ext, {}, Ext>
+          : {}
+        : {}
+    >
+  : never
+
+type ExtractStateExtensionsFromEnhancerTuple<
+  EnhancerTuple extends any[],
+  Acc extends {}
+> = EnhancerTuple extends [infer Head, ...infer Tail]
+  ? ExtractStateExtensionsFromEnhancerTuple<
+      Tail,
+      Acc &
+        (Head extends StoreEnhancer<any, infer StateExt>
+          ? IsAny<StateExt, {}, StateExt>
+          : {})
+    >
+  : Acc
+
+export type ExtractStateExtensions<E> = E extends EnhancerArray<
+  infer EnhancerTuple
+>
+  ? ExtractStateExtensionsFromEnhancerTuple<EnhancerTuple, {}>
+  : E extends ReadonlyArray<StoreEnhancer>
+  ? UnionToIntersection<
+      E[number] extends StoreEnhancer<any, infer StateExt>
+        ? StateExt extends {}
+          ? IsAny<StateExt, {}, StateExt>
           : {}
         : {}
     >