feat: Add the ability to type StoreEnhancers

packages/toolkit/src/configureStore.ts

@@ -20,7 +20,11 @@ import type {
   CurriedGetDefaultMiddleware,
 } from './getDefaultMiddleware'
 import { curryGetDefaultMiddleware } from './getDefaultMiddleware'
-import type { NoInfer, ExtractDispatchExtensions } from './tsHelpers'
+import type {
+  NoInfer,
+  ExtractDispatchExtensions,
+  ExtractStoreExtensions,
+} from './tsHelpers'
 
 const IS_PRODUCTION = process.env.NODE_ENV === 'production'
 
@@ -29,9 +33,9 @@ const IS_PRODUCTION = process.env.NODE_ENV === 'production'
  *
  * @public
  */
-export type ConfigureEnhancersCallback = (
-  defaultEnhancers: readonly StoreEnhancer[]
-) => StoreEnhancer[]
+export type ConfigureEnhancersCallback<E extends Enhancers = Enhancers> = (
+    defaultEnhancers: readonly StoreEnhancer[]
+) => [...E]
 
 /**
  * Options for `configureStore()`.
@@ -41,7 +45,8 @@ export type ConfigureEnhancersCallback = (
 export interface ConfigureStoreOptions<
   S = any,
   A extends Action = AnyAction,
-  M extends Middlewares<S> = Middlewares<S>
+  M extends Middlewares<S> = Middlewares<S>,
+  E extends Enhancers = Enhancers
 > {
   /**
    * A single reducer function that will be used as the root reducer, or an
@@ -52,7 +57,7 @@ export interface ConfigureStoreOptions<
   /**
    * An array of Redux middleware to install. If not supplied, defaults to
    * the set of middleware returned by `getDefaultMiddleware()`.
-   * 
+   *
    * @example `middleware: (gDM) => gDM().concat(logger, apiMiddleware, yourCustomMiddleware)`
    * @see https://redux-toolkit.js.org/api/getDefaultMiddleware#intended-usage
    */
@@ -79,7 +84,7 @@ export interface ConfigureStoreOptions<
   - if it is not, there could be two cases:
     - `ReducersMapObject<S, A>` is being passed in. In this case, we will call `combineReducers` on it and `CombinedState<S>` is correct
     - `Reducer<S, A>` is being passed in. In this case, actually `CombinedState<S>` is wrong and `S` would be correct.
-    As we cannot distinguish between those two cases without adding another generic paramter, 
+    As we cannot distinguish between those two cases without adding another generic parameter,
     we just make the pragmatic assumption that the latter almost never happens.
   */
   preloadedState?: PreloadedState<CombinedState<NoInfer<S>>>
@@ -92,21 +97,17 @@ export interface ConfigureStoreOptions<
    * and should return a new array (such as `[applyMiddleware, offline]`).
    * If you only need to add middleware, you can use the `middleware` parameter instead.
    */
-  enhancers?: StoreEnhancer[] | ConfigureEnhancersCallback
+  enhancers?: E | ConfigureEnhancersCallback<E>
 }
 
 type Middlewares<S> = ReadonlyArray<Middleware<{}, S>>
 
-/**
- * A Redux store returned by `configureStore()`. Supports dispatching
- * side-effectful _thunks_ in addition to plain actions.
- *
- * @public
- */
-export interface EnhancedStore<
+type Enhancers = ReadonlyArray<StoreEnhancer>
+
+interface ToolkitStore<
   S = any,
   A extends Action = AnyAction,
-  M extends Middlewares<S> = Middlewares<S>
+  M extends Middlewares<S> = Middlewares<S>,
 > extends Store<S, A> {
   /**
    * The `dispatch` method of your store, enhanced by all its middlewares.
@@ -116,19 +117,33 @@ export interface EnhancedStore<
   dispatch: ExtractDispatchExtensions<M> & Dispatch<A>
 }
 
+/**
+ * A Redux store returned by `configureStore()`. Supports dispatching
+ * side-effectful _thunks_ in addition to plain actions.
+ *
+ * @public
+ */
+export type EnhancedStore<
+  S = any,
+  A extends Action = AnyAction,
+  M extends Middlewares<S> = Middlewares<S>,
+  E extends Enhancers = Enhancers
+> = ToolkitStore<S, A, M> & ExtractStoreExtensions<E>
+
 /**
  * A friendly abstraction over the standard Redux `createStore()` function.
  *
- * @param config The store configuration.
+ * @param options The store configuration.
  * @returns A configured Redux store.
  *
  * @public
  */
 export function configureStore<
   S = any,
   A extends Action = AnyAction,
-  M extends Middlewares<S> = [ThunkMiddlewareFor<S>]
->(options: ConfigureStoreOptions<S, A, M>): EnhancedStore<S, A, M> {
+  M extends Middlewares<S> = [ThunkMiddlewareFor<S>],
+  E extends Enhancers = [StoreEnhancer]
+>(options: ConfigureStoreOptions<S, A, M, E>): EnhancedStore<S, A, M, E> {
   const curriedGetDefaultMiddleware = curryGetDefaultMiddleware<S>()
 
   const {
@@ -170,7 +185,7 @@ export function configureStore<
     )
   }
 
-  const middlewareEnhancer = applyMiddleware(...finalMiddleware)
+  const middlewareEnhancer: StoreEnhancer = applyMiddleware(...finalMiddleware)
 
   let finalCompose = compose
 
@@ -182,15 +197,15 @@ export function configureStore<
     })
   }
 
-  let storeEnhancers: StoreEnhancer[] = [middlewareEnhancer]
+  let storeEnhancers: Enhancers = [middlewareEnhancer]
 
   if (Array.isArray(enhancers)) {
     storeEnhancers = [middlewareEnhancer, ...enhancers]
   } else if (typeof enhancers === 'function') {
     storeEnhancers = enhancers(storeEnhancers)
   }
 
-  const composedEnhancer = finalCompose(...storeEnhancers) as any
+  const composedEnhancer = finalCompose(...storeEnhancers) as StoreEnhancer<any>
 
   return createStore(rootReducer, preloadedState, composedEnhancer)
 }

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

@@ -6,9 +6,10 @@ import type {
   Reducer,
   Store,
   Action,
+  StoreEnhancer
 } from 'redux'
 import { applyMiddleware } from 'redux'
-import type { PayloadAction, MiddlewareArray } from '@reduxjs/toolkit'
+import type { PayloadAction } from '@reduxjs/toolkit'
 import {
   configureStore,
   getDefaultMiddleware,
@@ -17,7 +18,6 @@ import {
 import type { ThunkMiddleware, ThunkAction, ThunkDispatch } from 'redux-thunk'
 import thunk from 'redux-thunk'
 import { expectNotAny, expectType } from './helpers'
-import type { IsAny, ExtractDispatchExtensions } from '../tsHelpers'
 
 const _anyMiddleware: any = () => () => () => {}
 
@@ -140,16 +140,53 @@ const _anyMiddleware: any = () => () => () => {}
  * Test: configureStore() accepts store enhancer.
  */
 {
-  configureStore({
-    reducer: () => 0,
-    enhancers: [applyMiddleware((store) => (next) => next)],
-  })
+  {
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: [applyMiddleware(() => next => next)]
+    })
+
+    expectType<Dispatch & ThunkDispatch<number, undefined, AnyAction>>(store.dispatch)
+  }
 
   configureStore({
     reducer: () => 0,
     // @ts-expect-error
     enhancers: ['not a store enhancer'],
   })
+
+  {
+    type SomePropertyStoreEnhancer = StoreEnhancer<{ someProperty: string }>
+
+    const somePropertyStoreEnhancer: SomePropertyStoreEnhancer = next => {
+      return (reducer, preloadedState) => {
+        return {
+          ...next(reducer, preloadedState),
+          someProperty: 'some value',
+        }
+      }
+    }
+
+    type AnotherPropertyStoreEnhancer = StoreEnhancer<{ anotherProperty: number }>
+
+    const anotherPropertyStoreEnhancer: AnotherPropertyStoreEnhancer = next => {
+      return (reducer, preloadedState) => {
+        return {
+          ...next(reducer, preloadedState),
+          anotherProperty: 123,
+        }
+      }
+    }
+
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: [somePropertyStoreEnhancer, anotherPropertyStoreEnhancer],
+    })
+
+    expectType<Dispatch & ThunkDispatch<number, undefined, AnyAction>>(store.dispatch)
+    expectType<string>(store.someProperty)
+    expectType<number>(store.anotherProperty)
+  }
 }
 
 /**

packages/toolkit/src/tsHelpers.ts

@@ -1,4 +1,4 @@
-import type { Middleware } from 'redux'
+import type { Middleware, StoreEnhancer } from 'redux'
 import type { MiddlewareArray } from './utils'
 
 /**
@@ -66,6 +66,15 @@ export type IsUnknownOrNonInferrable<T, True, False> = AtLeastTS35<
   IsEmptyObj<T, True, IsUnknown<T, True, False>>
 >
 
+/**
+ * Convert a Union type `(A|B)` to an intersection type `(A&B)`
+ */
+export type UnionToIntersection<U> = (
+  U extends any ? (k: U) => void : never
+) extends (k: infer I) => void
+  ? I
+  : never
+
 // Appears to have a convenient side effect of ignoring `never` even if that's not what you specified
 export type ExcludeFromTuple<T, E, Acc extends unknown[] = []> = T extends [
   infer Head,
@@ -80,7 +89,7 @@ type ExtractDispatchFromMiddlewareTuple<
 > = MiddlewareTuple extends [infer Head, ...infer Tail]
   ? ExtractDispatchFromMiddlewareTuple<
       Tail,
-      Acc & (Head extends Middleware<infer D, any> ? IsAny<D, {}, D> : {})
+      Acc & (Head extends Middleware<infer D> ? IsAny<D, {}, D> : {})
     >
   : Acc
 
@@ -92,14 +101,9 @@ export type ExtractDispatchExtensions<M> = M extends MiddlewareArray<
   ? ExtractDispatchFromMiddlewareTuple<[...M], {}>
   : never
 
-/**
- * Convert a Union type `(A|B)` to an intersection type `(A&B)`
- */
-export type UnionToIntersection<U> = (
-  U extends any ? (k: U) => void : never
-) extends (k: infer I) => void
-  ? I
-  : never
+export type ExtractStoreExtensions<E> = E extends any[]
+  ? UnionToIntersection<E[number] extends StoreEnhancer<infer Ext> ? Ext extends {} ? Ext : {} : {}>
+  : {}
 
 /**
  * Helper type. Passes T out again, but boxes it in a way that it cannot