Include case reducers in createSlice result (#209)

* Fix missing name field in type test

* Include provided case reducers in createSlice output

- Added `caseReducers` field to the createSlice return object
- Added test and type test for returned case reducer
- Removed "Enhanced" terminology from types, and replaced with
  "CaseReducerWithPrepare" and "CaseReducerDefinition"
- Restructured logic to only loop over reducer names once, and
  check case reducer definition type once per entry

* Add type tests for correct case reducer export inference

* Rewrite caseReducers types for correct inference of state + action

* Add type test for reducers with prepare callback

* Fix type inference for actions from reducers w/ prepare callbacks

* Clean up type names and usages

* Use a generic State type in ActionForReducer

* Re-switch Slice generics order

Author: markerikson

src/createSlice.test.ts

@@ -35,7 +35,7 @@ describe('createSlice', () => {
   })
 
   describe('when passing slice', () => {
-    const { actions, reducer } = createSlice({
+    const { actions, reducer, caseReducers } = createSlice({
       reducers: {
         increment: state => state + 1
       },
@@ -57,6 +57,12 @@ describe('createSlice', () => {
     it('should return the correct value from reducer', () => {
       expect(reducer(undefined, actions.increment())).toEqual(1)
     })
+
+    it('should include the generated case reducers', () => {
+      expect(caseReducers).toBeTruthy()
+      expect(caseReducers.increment).toBeTruthy()
+      expect(typeof caseReducers.increment).toBe('function')
+    })
   })
 
   describe('when mutating state object', () => {

src/createSlice.ts

@@ -18,7 +18,9 @@ export type SliceActionCreator<P> = PayloadActionCreator<P>
 
 export interface Slice<
   State = any,
-  ActionCreators extends { [key: string]: any } = { [key: string]: any }
+  CaseReducers extends SliceCaseReducerDefinitions<State, PayloadActions> = {
+    [key: string]: any
+  }
 > {
   /**
    * The slice name.
@@ -34,15 +36,20 @@ export interface Slice<
    * Action creators for the types of actions that are handled by the slice
    * reducer.
    */
-  actions: ActionCreators
+  actions: CaseReducerActions<CaseReducers>
+
+  caseReducers: SliceDefinedCaseReducers<CaseReducers, State>
 }
 
 /**
  * Options for `createSlice()`.
  */
 export interface CreateSliceOptions<
   State = any,
-  CR extends SliceCaseReducers<State, any> = SliceCaseReducers<State, any>
+  CR extends SliceCaseReducerDefinitions<
+    State,
+    any
+  > = SliceCaseReducerDefinitions<State, any>
 > {
   /**
    * The slice's name. Used to namespace the generated action types.
@@ -74,23 +81,23 @@ type PayloadActions<Types extends keyof any = string> = Record<
   PayloadAction
 >
 
-type EnhancedCaseReducer<State, Action extends PayloadAction> = {
+type CaseReducerWithPrepare<State, Action extends PayloadAction> = {
   reducer: CaseReducer<State, Action>
   prepare: PrepareAction<Action['payload']>
 }
 
-type SliceCaseReducers<State, PA extends PayloadActions> = {
+type SliceCaseReducerDefinitions<State, PA extends PayloadActions> = {
   [ActionType in keyof PA]:
     | CaseReducer<State, PA[ActionType]>
-    | EnhancedCaseReducer<State, PA[ActionType]>
+    | CaseReducerWithPrepare<State, PA[ActionType]>
 }
 
 type IfIsReducerFunctionWithoutAction<R, True, False = never> = R extends (
   state: any
 ) => any
   ? True
   : False
-type IfIsEnhancedReducer<R, True, False = never> = R extends {
+type IfIsCaseReducerWithPrepare<R, True, False = never> = R extends {
   prepare: Function
 }
   ? True
@@ -106,8 +113,21 @@ type PrepareActionForReducer<R> = R extends { prepare: infer Prepare }
   ? Prepare
   : never
 
-type CaseReducerActions<CaseReducers extends SliceCaseReducers<any, any>> = {
-  [Type in keyof CaseReducers]: IfIsEnhancedReducer<
+type ActionForReducer<R, S> = R extends (
+  state: S,
+  action: PayloadAction<infer P>
+) => S
+  ? PayloadAction<P>
+  : R extends {
+      reducer(state: any, action: PayloadAction<infer P>): any
+    }
+  ? PayloadAction<P>
+  : unknown
+
+type CaseReducerActions<
+  CaseReducers extends SliceCaseReducerDefinitions<any, any>
+> = {
+  [Type in keyof CaseReducers]: IfIsCaseReducerWithPrepare<
     CaseReducers[Type],
     ActionCreatorWithPreparedPayload<
       PrepareActionForReducer<CaseReducers[Type]>
@@ -122,6 +142,16 @@ type CaseReducerActions<CaseReducers extends SliceCaseReducers<any, any>> = {
   >
 }
 
+type SliceDefinedCaseReducers<
+  CaseReducers extends SliceCaseReducerDefinitions<any, any>,
+  State = any
+> = {
+  [Type in keyof CaseReducers]: CaseReducer<
+    State,
+    ActionForReducer<CaseReducers[Type], State>
+  >
+}
+
 type NoInfer<T> = [T][T extends any ? 0 : never]
 
 type SliceCaseReducersCheck<S, ACR> = {
@@ -134,9 +164,9 @@ type SliceCaseReducersCheck<S, ACR> = {
     : {}
 }
 
-type RestrictEnhancedReducersToMatchReducerAndPrepare<
+type RestrictCaseReducerDefinitionsToMatchReducerAndPrepare<
   S,
-  CR extends SliceCaseReducers<S, any>
+  CR extends SliceCaseReducerDefinitions<S, any>
 > = { reducers: SliceCaseReducersCheck<S, NoInfer<CR>> }
 
 function getType(slice: string, actionKey: string): string {
@@ -153,54 +183,59 @@ function getType(slice: string, actionKey: string): string {
  */
 export function createSlice<
   State,
-  CaseReducers extends SliceCaseReducers<State, any>
+  CaseReducers extends SliceCaseReducerDefinitions<State, any>
 >(
   options: CreateSliceOptions<State, CaseReducers> &
-    RestrictEnhancedReducersToMatchReducerAndPrepare<State, CaseReducers>
-): Slice<State, CaseReducerActions<CaseReducers>>
+    RestrictCaseReducerDefinitionsToMatchReducerAndPrepare<State, CaseReducers>
+): Slice<State, CaseReducers>
 
 // internal definition is a little less restrictive
 export function createSlice<
   State,
-  CaseReducers extends SliceCaseReducers<State, any>
+  CaseReducers extends SliceCaseReducerDefinitions<State, any>
 >(
   options: CreateSliceOptions<State, CaseReducers>
-): Slice<State, CaseReducerActions<CaseReducers>> {
+): Slice<State, CaseReducers> {
   const { name, initialState } = options
   if (!name) {
     throw new Error('`name` is a required option for createSlice')
   }
   const reducers = options.reducers || {}
   const extraReducers = options.extraReducers || {}
-  const actionKeys = Object.keys(reducers)
-
-  const reducerMap = actionKeys.reduce((map, actionKey) => {
-    let maybeEnhancedReducer = reducers[actionKey]
-    map[getType(name, actionKey)] =
-      typeof maybeEnhancedReducer === 'function'
-        ? maybeEnhancedReducer
-        : maybeEnhancedReducer.reducer
-    return map
-  }, extraReducers)
-
-  const reducer = createReducer(initialState, reducerMap)
-
-  const actionMap = actionKeys.reduce(
-    (map, action) => {
-      let maybeEnhancedReducer = reducers[action]
-      const type = getType(name, action)
-      map[action] =
-        typeof maybeEnhancedReducer === 'function'
-          ? createAction(type)
-          : createAction(type, maybeEnhancedReducer.prepare)
-      return map
-    },
-    {} as any
-  )
+  const reducerNames = Object.keys(reducers)
+
+  const sliceCaseReducersByName: Record<string, CaseReducer> = {}
+  const sliceCaseReducersByType: Record<string, CaseReducer> = {}
+  const actionCreators: Record<string, PayloadActionCreator> = {}
+
+  reducerNames.forEach(reducerName => {
+    const maybeReducerWithPrepare = reducers[reducerName]
+    const type = getType(name, reducerName)
+
+    let caseReducer: CaseReducer<State, any>
+    let prepareCallback: PrepareAction<any> | undefined
+
+    if (typeof maybeReducerWithPrepare === 'function') {
+      caseReducer = maybeReducerWithPrepare
+    } else {
+      caseReducer = maybeReducerWithPrepare.reducer
+      prepareCallback = maybeReducerWithPrepare.prepare
+    }
+
+    sliceCaseReducersByName[reducerName] = caseReducer
+    sliceCaseReducersByType[type] = caseReducer
+    actionCreators[reducerName] = prepareCallback
+      ? createAction(type, prepareCallback)
+      : createAction(type)
+  })
+
+  const finalCaseReducers = { ...extraReducers, ...sliceCaseReducersByType }
+  const reducer = createReducer(initialState, finalCaseReducers)
 
   return {
     name,
     reducer,
-    actions: actionMap
+    actions: actionCreators as any,
+    caseReducers: sliceCaseReducersByName as any
   }
 }

type-tests/files/createSlice.typetest.ts

@@ -147,6 +147,57 @@ function expectType<T>(t: T) {
   expectType<string>(counter.actions.concatMetaStrLen('test').meta)
 }
 
+/*
+ * Test: returned case reducer has the correct type
+ */
+{
+  const counter = createSlice({
+    name: 'counter',
+    initialState: 0,
+    reducers: {
+      increment(state, action: PayloadAction<number>) {
+        return state + action.payload
+      },
+      decrement: {
+        reducer(state, action: PayloadAction<number>) {
+          return state - action.payload
+        },
+        prepare(amount: number) {
+          return { payload: amount }
+        }
+      }
+    }
+  })
+
+  // Should match positively
+  expectType<(state: number, action: PayloadAction<number>) => number | void>(
+    counter.caseReducers.increment
+  )
+
+  // Should match positively for reducers with prepare callback
+  expectType<(state: number, action: PayloadAction<number>) => number | void>(
+    counter.caseReducers.decrement
+  )
+
+  // Should not mismatch the payload if it's a simple reducer
+  // typings:expect-error
+  expectType<(state: number, action: PayloadAction<string>) => number | void>(
+    counter.caseReducers.increment
+  )
+
+  // Should not mismatch the payload if it's a reducer with a prepare callback
+  // typings:expect-error
+  expectType<(state: number, action: PayloadAction<string>) => number | void>(
+    counter.caseReducers.decrement
+  )
+
+  // Should not include entries that don't exist
+  // typings:expect-error
+  expectType<(state: number, action: PayloadAction<string>) => number | void>(
+    counter.caseReducers.someThingNonExistant
+  )
+}
+
 /*
  * Test: prepared payload does not match action payload - should cause an error.
  */
@@ -180,6 +231,7 @@ function expectType<T>(t: T) {
   }
 
   const mySlice = createSlice({
+    name: 'name',
     initialState,
     reducers: {
       setName: (state, action) => {