How would I create a Zustand store selector hook wrapper in Typescript

I am trying to use Zustand as a replacement for React Context to provide a shared state amongst a subset of components, and one of the patterns I find myself using a lot is the shallow, multi-property selector, like so:

// store.ts
interface State {
  data: object,
  setData: (newData: State['data']) => void,
}

const useStore = create<State>(set => ({
  data: {},
  setData: (data) => set({ data }),
});


// Component.tsx
import shallow from 'zustand/shallow';
import useStore from './store.ts';

const { data, setData } = useStore(
  ({ data, setData }) => ({ data, setData }),
  shallow
);

// Use 'data' and 'setData'

My understanding is that I need to use the shallow equality check when doing this sort of multi-property selector, but this is clearly a very verbose pattern to write out the keys multiple times and always add the shallow option, especially if it's more than just one or two keys.

I'd like to create a wrapper hook for useStore that will just take a list of keys and handle the rest of this for me, like this:

// store.ts
import pick from 'lodash/pick';

const useStoreState = (keys: (keyof State)[]) =>
  useStore(state => pick(state, keys), shallow)

// Component.tsx
import { useStoreState } from './store';

// Ideally the return value here has the same type as the example above
const { data, setData } = useStoreState(['data', 'setData']);

This functions just fine (at least I think it does - I don't know for sure if the pick breaks the selector subscription stuff), but I can't for the life of me get the typing to work out correctly.

Zustand's types are pretty complicated and do a lot of inference, but does someone know a way to accomplish this? Am I missing another approach to address this issue?

Cheers!