How to modularize this react state container?

user_498624 0
Viewed 98

So at work we have this awesome state container hook we built to use in our React application and associated packages. First a little background on this hook and what I'd like to preserve before getting to what I want to do with it. Here's the working code. You'll notice it's commented for easy copy and paste to create new ones.

// Set this accordingly.
const hookName = "MyState";

// Default state values. Make sure your values have explicit types set.
const initialState = {
  nums: [] as Number[]
};

// Available actions and the data they require. Specify "null" if action requires no data.
type Actions = {
  RESET_NUMS: null,
  ADD_NUM: Number,
  SET_NUMS: Number[]
};

// Action handler methods. There must be a handler for each action listed above.
const actionHandlers: ActionHandlers = {
  RESET_NUMS: () => ({ nums: [] }),
  ADD_NUM: ({ nums }, num) => {
    nums.push(num);
    return { nums };
  },
  SET_NUMS: ({}, nums) => ({ nums })
};

// Rename these exports accordingly.
export const MyStateProvider = Provider;
export const useMyState = useContextState;



/**************************************************
 *                                                *
 * You do not need to modify anything below this. *
 *                                                *
 **************************************************/

/* Context Hook Boilerplate */

import React, { useReducer, useContext, ReactNode } from "react";
import _cloneDeep from "lodash/cloneDeep";

const Context = React.createContext<
    {
        state: State;
        dispatch: Dispatch
    } | undefined
>(undefined);

function reducer(state: State, action: Action): State {
    const stateClone = _cloneDeep(state);
    const newState = actionHandlers[action.type](stateClone, action.data as any);
    if (!newState) return state;
    return { ...stateClone, ...newState };
}

function Provider({children, defaultState}: {children: ReactNode, defaultState?: State}) {
    const [state, reducerDispatch] = useReducer(reducer, defaultState ?? initialState);
    const dispatch: Dispatch = (type, data) => reducerDispatch({type, data} as Action);
    return (
        <Context.Provider value={{ state, dispatch }}>
            {children}
        </Context.Provider>
    );
}

function useContextState() {
    const context = useContext(Context);
    if (context === undefined) {
        throw new Error(`use${hookName} must be used within a ${hookName}Provider`);
    }
    return context;
}

/* TypeScript Boilerplate */

type Dispatch = <A extends Actions, T extends keyof A, D extends A[T]>(type: T, ...data: (D extends null ? [] : [D])) => void
type State = typeof initialState;
type ActionsMap = {
  [K in keyof Actions]: Actions[K] extends null ? { type: K, data?: undefined } : { type: K, data: Actions[K] }
};
type Action = ActionsMap[keyof Actions];
type ActionHandlers = {
    [K in keyof Actions]: Actions[K] extends null ? (state: State) => Partial<State> | void : (state: State, data: Actions[K]) => Partial<State> | void;
};

What's cool about it is that the state and dispatch functions are very rigidly typed.

  • The type of initialState is inferred from its structure so when copying this code to make a new state container hook you don't have to write a State type and then also a default state object.
  • The Actions type defines actions that are available to be dispatched and the type of each key is the type that is used to enforce the second argument to our dispatch function.
  • The actionHandlers object must have a handler for each action defined on Actions and the handler functions have to conform exactly, receiving the state object as the first argument, and the second argument being a data object that matches the relevant type value on Actions.
  • The dispatch function that consumers of this hook use is almost impossible to misuse. The first argument only accepts valid actions from Actions keys and the second argument is either not allowed in the case of the Actions type being null or is forced to conform to the Actions type specified.

Here's a typescript playground if you wanna mess with it (go to line 65 to tinker with making dispatch calls)

Okay so what you might notice is that the bottom half of the code snippet is boilerplate. Well we use these hooks a lot so we are copying and pasting them into various modules in our project. It works great but it's not very DRY and any change to the boilerplate obviously has to be propagated to every single hook. Tedious.

I want to modularize the boilerplate bit but as I start trying to think about it my mind gets boggled about how to maintain the rigid type structures behind the scenes. Here's a brainstorm of what it might be like to consume such a module:

import StateContainer from "@company/react-state-container";

const myState = new StateContainer("MyState");

myState.defaultState({
  nums: [] as number[]
});

myState.actionHandler(
  "RESET_NUMS",
  () => ({ nums: [] })
);

myState.actionHandler(
  "ADD_NUM",
  ({ nums }, num: number) => ({ nums: [...nums, num] }) 
);

myState.actionHandler(
  "SET_NUMS",
  ({}, nums: number[]) => ({ nums })
);

const {
  Provider: MyStateProvider,
  useStateContainer: useMyState
} = myState;

export { MyStateProvider, useMyState };

This is just the first thing that came to my head as far as an API goes. Maybe you can think of something better. While consuming the API makes sense to me, how to write this module is where I'm tripping over myself. For one example: How can I infer the type of the default state like I do in the original code? I can write the defaultState class method using generics but then how do I propagate that generic type into the rest of the class outside of that one method? Is that even possible?

Another question I had pretty quickly. Should I expect the user to provide user-defined types, like for Actions? Should the consumer have to pass in TypeScript stuff at all or is there a way I can just have them pass only the imperative code and then I can layer on some type magic for those who are taking advantage of TypeScript?

The more I start trying to modularize that code the more I feel like I'm just taking the wrong approach altogether or maybe what I want to do is a fool's errand in the first place.

PS - I know I know, Redux. Sadly I don't get to decide.

react-hooks typescript reactjs javascript
1 Answers

PS - I know I know, Redux. Sadly I don't get to decide.

Yes, you're basically re-creating Redux here. More specifically, you're trying to re-create the createSlice functionality of Redux Toolkit. You want to define a mapping of action names to what the action does, and then have the reducer get created automatically. So we can use that prior art to get an idea of how this might work.

Your current brainstorm involves calling functions on the StateContainer object after it has been created. Those functions need to change the types of the StateContainer. This is doable, but it's easier to create the object in one go with all of the information up front.

Let's think about what information needs to be provided and what information needs to be returned.

We need a name, an initialState, and a bunch of actions:

type Config<S, AH> = {
  name: string;
  initialState: S;
  actionHandlers: AH;
}

The type of the state and the type of the actions are generics where S represents State and AH represents ActionHandlers. I'm using letters to make it clear what's a generic and what's an actual type.

We want to put some sort of constraint on the actions. It should be an object whose keys are strings (the action names) and whose values are functions. Those functions take the state and a payload (which will have a different type for each action) and return a new state. Actually your code says that we return Partial<State> | void. I'm not sure what that void accomplishes? But we get this:

type GenericActionHandler<S, P> = (state: S, payload: P) => Partial<S> | void;

type Config<S, AH extends Record<string, GenericActionHandler<S, any>>> = {
...

Our utility is going to take that Config and return a Provider and a hook with properties state and dispatch. It's the dispatch that requires us to do the fancy TypeScript inference in order to get the correct types for the type and data arguments. FYI, having those as two separate arguments does make is slightly harder to ensure that you've got a matching pair.

The typing is similar to the "TypeScript Boilerplate" that you had before. The main difference is that we are working backwards from ActionHandlers to Actions. The nested ternary here handles the situation where there is no second argument.

type Actions<AH> = {
  [K in keyof AH]: AH[K] extends GenericActionHandler<any, infer P>
    ? (unknown extends P ? never : P )
    : never;
};

type TypePayloadPair<AH> = {
  [K in keyof AH]: Actions<AH>[K] extends null | undefined
    ? [K]
    : [K, Actions<AH>[K]];
}[keyof AH];

type Dispatch<AH> = (...args: TypePayloadPair<AH>) => void;

So now, finally, we know the return type of the StateContainer object. Given a state type S and an action handlers object AH, the container type is:

type StateContainer<S, AH> = {
  Provider: React.FC<{defaultState?: S}>;
  useContextState: () => {
    state: S;
    dispatch: Dispatch<AH>;
  }
}

We'll use that as the return type for the factory function, which I am calling createStateContainer. The argument type is the Config that we wrote earlier.

The conversion of type and data to {type, data} as Action is not really necessary because the React useReducer hook doesn't make any requirements about the action type. You can avoid all as assertions within your function if you pass along the pair of arguments from dispatch as-is.

export default function createStateContainer<
    S,
    AH extends Record<string, GenericActionHandler<S, unknown>>
>({
    name,
    initialState,
    actionHandlers,
}: Config<S, AH>): StateContainer<S, AH> {

  const Context = React.createContext<
      {
          state: S;
          dispatch: Dispatch<AH>
      } | undefined
  >(undefined);

  function reducer(state: S, [type, payload]: TypePayloadPair<AH>): S {
      const stateClone = _cloneDeep(state);
      const newState = actionHandlers[type](stateClone, payload);
      if (!newState) return state;
      return { ...stateClone, ...newState };
  }

  function Provider({children, defaultState}: {children?: ReactNode, defaultState?: S}) {
      const [state, reducerDispatch] = useReducer(reducer, defaultState ?? initialState);
      const dispatch: Dispatch<AH> = (...args) => reducerDispatch(args);
      
      return (
          <Context.Provider value={{ state, dispatch }}>
              {children}
          </Context.Provider>
      );
  }

  function useContextState() {
      const context = useContext(Context);
      if (context === undefined) {
          throw new Error(`use${name} must be used within a ${name}Provider`);
      }
      return context;
  }

  return {
    Provider,
    useContextState
  }
}

Creating an instance has gotten much, much simpler:

import createStateContainer from "./StateContainer";

export const {
  Provider: MyStateProvider,
  useContextState: useMyState
} = createStateContainer({
  name: "MyState",
  initialState: {
    nums: [] as number[]
  },
  actionHandlers: {
    RESET_NUMS: () => ({ nums: [] }),
    ADD_NUM: ({ nums }, num: number) => {
      nums.push(num);
      return { nums };
    },
    SET_NUMS: ({}, nums: number[]) => ({ nums })
  }
});

Complete Code and Demo

user_104315740
Related