typed useReducer for react context

CreateContextReducer.tsx

import type { Dispatch } from 'react';
import { useReducer } from 'react';

type InitialValue = {
  search: {
    value?: string;
  };
  results: {
    loading: boolean;
    error: boolean;
    data: { id: number }[] | null;
  };
  pagination: {
    limit: number;
    offset: number;
  };
};

const initialValue = {
  search: {
    value: '',
  },
  results: {
    loading: false,
    error: false,
    data: null,
  },
  pagination: {
    limit: 25,
    offset: 0,
  },
};

type NullIdentity<T> = T extends null ? any : T;

type Compose<S, P> = P extends object
  ? {
      [T in keyof P]: T extends keyof S
        ? keyof Partial<P[T]> extends keyof Partial<S[T]>
          ? Compose<S[T], P[T]>
          : S[T] extends object
            ? Partial<S[T]>
            : S[T] extends null
              ? any
              : S[T]
        : NullIdentity<P[T]>;
    }
  : NullIdentity<P>;

export type Action<S, H, T = keyof H> = T extends keyof H
  ? H[T] extends (...args: any) => any
    ? Parameters<H[T]> extends [S, infer X]
      ? X extends object
        ? { type: T; payload: Compose<S, X> }
        : { type: T; payload: X }
      : { type: T }
    : never
  : never;

export type Handler<S, P> = (state: S, payload?: P) => S;

export type Handlers<S> = {
  [K: string]: Handler<S, unknown>;
};

export function createContextReducer<S = unknown, H = Handlers<S>>(
  handlers: H,
  initialState: S,
) {
  const reducer = (state: S, action: Action<S, H>): S => {
    const handler = handlers[action.type] as Handler<S, unknown> | undefined;
    if (handler) {
      if ('payload' in action) {
        return handler(state, action.payload);
      }
      return handler(state);
    }
    return state;
  };

  const context = {
    state: initialState,
    dispatch: (() => undefined) as Dispatch<[Action<S, H>]>,
  };

  return [reducer, context] as const;
}

export function jsContextReducer(handlers, initialState) {
  const reducer = (state, action) => {
    const handler = handlers?.[action?.type];
    if (handler) {
      return handler(state, action?.payload);
    }
    return state;
  };

  const context = {
    state: initialState,
    dispatch: () => undefined,
  };

  return [reducer, context];
}

const [reducer, context] = createContextReducer(
  {
    SET_RESULTS_DATA: (state, { results: { data } }) => ({
      ...state,
      results: {
        loading: false,
        data,
        error: false,
      },
    }),
    // payload is partial of initialValue. should enforce `value` as a string
    SET_SEARCH: (state, { search: { value } }) => ({
      ...state,
      search: { value },
    }),
    SET_RESULTS: (state, { results }) => ({
      ...state,
      results: {
        ...state.results,
        ...results,
      },
    }),
    // payload is partial of initialValue. should enforce `loading` as a boolean
    SET_RESULTS_LOADING: (state, { results: { loading } }) => ({
      ...state,
      results: {
        loading,
        data: null,
        error: false,
      },
    }),
    // payload is not partial of `initialValue`. enforce based on specified type
    SET_RESULTS_ERROR: (state, { error }: { error: string }) => ({
      ...state,
      results: {
        loading: false,
        data: null,
        error,
      },
    }),
    SET_IDENTITY: (state) => state,
  },
  initialValue,
);

function Test() {
  const [state, dispatch] = useReducer(reducer, initialValue);
  // ts error on `payload`
  dispatch({ type: 'SET_IDENTITY', payload: { test: false } });
  // correct
  dispatch({ type: 'SET_IDENTITY' });

  // data should be `any` here since it's `null` in initialValue
  dispatch({
    type: 'SET_RESULTS_DATA',
    payload: { results: { data: 'thing' } },
  });
  dispatch({
    type: 'SET_RESULTS_DATA',
    payload: {
      results: { data: [{ id: 2 }, { id: 3 }], fake: true },
      fake: true,
    },
  });
  // ts error on `fake`
  dispatch({
    type: 'SET_RESULTS_LOADING',
    payload: { results: { loading: true, fake: 'thing' } },
  });

  // enforce partial of `results`
  dispatch({ type: 'SET_RESULTS', payload: { results: { error: 'string' } } });
  dispatch({ type: 'SET_RESULTS', payload: { results: { error: true } } });

  // should enforce `value` being a string
  dispatch({ type: 'SET_SEARCH', payload: { search: { value: 'test' } } });
  // ts error on `value` as `boolean`
  dispatch({ type: 'SET_SEARCH', payload: { search: { value: true } } });
  // `error` should be enforced as a `boolean` since its type is defined in the function
  dispatch({ type: 'SET_RESULTS_ERROR', payload: { error: 'error' } });
  dispatch({ type: 'SET_RESULTS_ERROR', payload: { error: false } });
}