Typescript types for a hook that returns a piece of the Redux state

I have this hook named useReduxState which I use to retrieve a piece of state from Redux by doing:

const STATE_A = useReduxState("STATE_A");

And I'm having trouble while adding Typescript to it.

These are the types that I'm working with:

type TYPE_A = {
  fooA: string,
  barA: string,
}

type TYPE_B = {
  fooB: string,
  barB: string
}

interface REDUX_STATE {
  STATE_A: TYPE_A,
  STATE_B: TYPE_B,
}

And this is the code for the hook (it uses the useSelector hook from the react-redux package);

function useReduxState(STATE_NAME: REDUX_STATE_KEYS): REDUX_STATE_VALUES {
  
  const STATE = useSelector((state: REDUX_STATE) : REDUX_STATE_VALUES => state[STATE_NAME]);
  return STATE;

}

My goal is to be able to write something as:

const { fooA, barA } = useReduxState("STATE_A");

And I want Typescript to know that I can access those two properties fooA and barA because they are present in TYPE_A which is the type for the STATE_A.

But I'm getting these errors:

Property 'fooA' does not exist on type 'REDUX_STATE_VALUES'.ts(2339)
Property 'barA' does not exist on type 'REDUX_STATE_VALUES'.ts(2339)

Here is the type REDUX_STATE_VALUES

It seems correct, but I want Typescript to know that when I pass "STATE_A" as an argument, the return type should be TYPE_A, so I can access its properties fooA and barA. And vice-versa for "STATE_B" and its type.

How can I achieve that? Should I be using generics in this case to be able to get this behavior?

you would need a generic so that each call to your function will determine which field it is accessing:

function useReduxState<K extends keyof REDUX_STATE>(STATE_NAME: K): REDUX_STATE[K] {

the argument will be one keyof REDUX_STATE and the return type will be the corresponding value, so when you are passed a specific string it will figure out the correct return type.

you will also need to remove the explicit return annotation on your function, it will infer the correct type from usage and you need more specific than REDUX_STATE_VALUES

useSelector((state: REDUX_STATE) => state[STATE_NAME])