create EnhancerArray class for proper inference of enhancer extensions

Author: ben.durrant

packages/toolkit/src/configureStore.ts

@@ -25,6 +25,7 @@ import type {
   ExtractDispatchExtensions,
   ExtractStoreExtensions,
 } from './tsHelpers'
+import { EnhancerArray } from './utils'
 
 const IS_PRODUCTION = process.env.NODE_ENV === 'production'
 
@@ -34,8 +35,8 @@ const IS_PRODUCTION = process.env.NODE_ENV === 'production'
  * @public
  */
 export type ConfigureEnhancersCallback<E extends Enhancers = Enhancers> = (
-    defaultEnhancers: readonly StoreEnhancer[]
-) => [...E]
+  defaultEnhancers: EnhancerArray<[StoreEnhancer<{}, {}>]>
+) => E
 
 /**
  * Options for `configureStore()`.
@@ -107,7 +108,7 @@ type Enhancers = ReadonlyArray<StoreEnhancer>
 export 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.
@@ -197,12 +198,13 @@ export function configureStore<
     })
   }
 
-  let storeEnhancers: Enhancers = [middlewareEnhancer]
+  const defaultEnhancers = new EnhancerArray(middlewareEnhancer)
+  let storeEnhancers: Enhancers = defaultEnhancers
 
   if (Array.isArray(enhancers)) {
     storeEnhancers = [middlewareEnhancer, ...enhancers]
   } else if (typeof enhancers === 'function') {
-    storeEnhancers = enhancers(storeEnhancers)
+    storeEnhancers = enhancers(defaultEnhancers)
   }
 
   const composedEnhancer = finalCompose(...storeEnhancers) as StoreEnhancer<any>

packages/toolkit/src/index.ts

@@ -103,7 +103,7 @@ export type {
   // types
   ActionReducerMapBuilder,
 } from './mapBuilders'
-export { MiddlewareArray } from './utils'
+export { MiddlewareArray, EnhancerArray } from './utils'
 
 export { createEntityAdapter } from './entities/create_adapter'
 export type {

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

@@ -0,0 +1,103 @@
+import { configureStore } from '@reduxjs/toolkit'
+import type { StoreEnhancer } from 'redux'
+
+declare const expectType: <T>(t: T) => T
+
+declare const enhancer1: StoreEnhancer<{
+  has1: true
+}>
+
+declare const enhancer2: StoreEnhancer<{
+  has2: true
+}>
+
+{
+  // prepend single element
+  {
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: (dE) => dE.prepend(enhancer1),
+    })
+    expectType<true>(store.has1)
+
+    // @ts-expect-error
+    expectType<true>(store.has2)
+  }
+
+  // prepend multiple (rest)
+  {
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: (dE) => dE.prepend(enhancer1, enhancer2),
+    })
+    expectType<true>(store.has1)
+    expectType<true>(store.has2)
+
+    // @ts-expect-error
+    expectType<true>(store.has3)
+  }
+
+  // prepend multiple (array notation)
+  {
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: (dE) => dE.prepend([enhancer1, enhancer2] as const),
+    })
+    expectType<true>(store.has1)
+    expectType<true>(store.has2)
+
+    // @ts-expect-error
+    expectType<true>(store.has3)
+  }
+
+  // concat single element
+  {
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: (dE) => dE.concat(enhancer1),
+    })
+    expectType<true>(store.has1)
+
+    // @ts-expect-error
+    expectType<true>(store.has2)
+  }
+
+  // prepend multiple (rest)
+  {
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: (dE) => dE.concat(enhancer1, enhancer2),
+    })
+    expectType<true>(store.has1)
+    expectType<true>(store.has2)
+
+    // @ts-expect-error
+    expectType<true>(store.has3)
+  }
+
+  // concat multiple (array notation)
+  {
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: (dE) => dE.concat([enhancer1, enhancer2] as const),
+    })
+    expectType<true>(store.has1)
+    expectType<true>(store.has2)
+
+    // @ts-expect-error
+    expectType<true>(store.has3)
+  }
+
+  // concat and prepend
+  {
+    const store = configureStore({
+      reducer: () => 0,
+      enhancers: (dE) => dE.concat(enhancer1).prepend(enhancer2),
+    })
+    expectType<true>(store.has1)
+    expectType<true>(store.has2)
+
+    // @ts-expect-error
+    expectType<true>(store.has3)
+  }
+}

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

@@ -231,9 +231,7 @@ describe('configureStore', () => {
 
       const store = configureStore({
         reducer,
-        enhancers: (defaultEnhancers) => {
-          return [...defaultEnhancers, dummyEnhancer]
-        },
+        enhancers: (defaultEnhancers) => defaultEnhancers.concat(dummyEnhancer),
       })
 
       expect(dummyEnhancerCalled).toBe(true)

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

@@ -194,6 +194,20 @@ const _anyMiddleware: any = () => () => () => {}
     )
     expectType<string>(store.someProperty)
     expectType<number>(store.anotherProperty)
+
+    const storeWithCallback = configureStore({
+      reducer: () => 0,
+      enhancers: (defaultEnhancers) =>
+        defaultEnhancers
+          .prepend(anotherPropertyStoreEnhancer)
+          .concat(somePropertyStoreEnhancer),
+    })
+
+    expectType<Dispatch & ThunkDispatch<number, undefined, AnyAction>>(
+      store.dispatch
+    )
+    expectType<string>(storeWithCallback.someProperty)
+    expectType<number>(storeWithCallback.anotherProperty)
   }
 }
 

packages/toolkit/src/tsHelpers.ts

@@ -1,5 +1,5 @@
 import type { Middleware, StoreEnhancer } from 'redux'
-import type { MiddlewareArray } from './utils'
+import type { EnhancerArray, MiddlewareArray } from './utils'
 
 /**
  * return True if T is `any`, otherwise return False
@@ -101,9 +101,29 @@ export type ExtractDispatchExtensions<M> = M extends MiddlewareArray<
   ? ExtractDispatchFromMiddlewareTuple<[...M], {}>
   : never
 
-export type ExtractStoreExtensions<E> = E extends any[]
-  ? UnionToIntersection<E[number] extends StoreEnhancer<infer Ext> ? Ext extends {} ? Ext : {} : {}>
-  : {}
+type ExtractStoreExtensionsFromEnhancerTuple<
+  EnhancerTuple extends any[],
+  Acc extends {}
+> = EnhancerTuple extends [infer Head, ...infer Tail]
+  ? ExtractStoreExtensionsFromEnhancerTuple<
+      Tail,
+      Acc & (Head extends StoreEnhancer<infer Ext> ? IsAny<Ext, {}, Ext> : {})
+    >
+  : Acc
+
+export type ExtractStoreExtensions<E> = E extends EnhancerArray<
+  infer EnhancerTuple
+>
+  ? ExtractStoreExtensionsFromEnhancerTuple<EnhancerTuple, {}>
+  : E extends ReadonlyArray<StoreEnhancer>
+  ? UnionToIntersection<
+      E[number] extends StoreEnhancer<infer Ext>
+        ? Ext extends {}
+          ? Ext
+          : {}
+        : {}
+    >
+  : never
 
 /**
  * Helper type. Passes T out again, but boxes it in a way that it cannot

packages/toolkit/src/utils.ts

@@ -1,5 +1,5 @@
 import createNextState, { isDraftable } from 'immer'
-import type { Middleware } from 'redux'
+import type { Middleware, StoreEnhancer } from 'redux'
 
 export function getTimeMeasureUtils(maxDelay: number, fnName: string) {
   let elapsed = 0
@@ -70,6 +70,49 @@ export class MiddlewareArray<
   }
 }
 
+/**
+ * @public
+ */
+export class EnhancerArray<
+  Enhancers extends StoreEnhancer<any, any>[]
+> extends Array<Enhancers[number]> {
+  constructor(...items: Enhancers)
+  constructor(...args: any[]) {
+    super(...args)
+    Object.setPrototypeOf(this, EnhancerArray.prototype)
+  }
+
+  static get [Symbol.species]() {
+    return EnhancerArray as any
+  }
+
+  concat<AdditionalEnhancers extends ReadonlyArray<StoreEnhancer<any, any>>>(
+    items: AdditionalEnhancers
+  ): EnhancerArray<[...Enhancers, ...AdditionalEnhancers]>
+
+  concat<AdditionalEnhancers extends ReadonlyArray<StoreEnhancer<any, any>>>(
+    ...items: AdditionalEnhancers
+  ): EnhancerArray<[...Enhancers, ...AdditionalEnhancers]>
+  concat(...arr: any[]) {
+    return super.concat.apply(this, arr)
+  }
+
+  prepend<AdditionalEnhancers extends ReadonlyArray<StoreEnhancer<any, any>>>(
+    items: AdditionalEnhancers
+  ): EnhancerArray<[...AdditionalEnhancers, ...Enhancers]>
+
+  prepend<AdditionalEnhancers extends ReadonlyArray<StoreEnhancer<any, any>>>(
+    ...items: AdditionalEnhancers
+  ): EnhancerArray<[...AdditionalEnhancers, ...Enhancers]>
+
+  prepend(...arr: any[]) {
+    if (arr.length === 1 && Array.isArray(arr[0])) {
+      return new EnhancerArray(...arr[0].concat(this))
+    }
+    return new EnhancerArray(...arr.concat(this))
+  }
+}
+
 export function freezeDraftable<T>(val: T) {
   return isDraftable(val) ? createNextState(val, () => {}) : val
 }