ImmerState function with types

immerState.ts

import React from "react";
import produce from "immer";

// get type of partial of an object with type T
export type IPartialObject<T> = { [K in keyof T]?: T[K] };

export interface ImmerStateFunc<S> {
  (f: IPartialObject<S> | ((s: S) => void), callback?: () => void): void;
}

export let immerHelpers = {
  setState: (...f: any[]) => void {},
  immerState: function<S>(f: IPartialObject<S> | ((s: S) => void), callback?: () => void) {
    if (typeof f === "function") {
      this.setState(produce<any>(f as (s: S) => void), callback);
    } else {
      let partialState = f;
      let newState = produce<S>((state) => {
        Object.assign(state, partialState);
      });
      this.setState(newState, callback);
    }
  },
};

interface IState {
  isLoading: boolean;
  isLoaded: boolean;
}

// actual usage

class A extends React.Component<any, IState> {
  render() {
    return null;
  }

  immerState = immerHelpers.immerState as ImmerStateFunc<IState>;

  onchange() {
    this.immerState({
      isLoading: true,
    });
  }
}