typing error when union of atoms are used

[double-thinker]

Bug Description

Using union of atoms in get and set tirggers a typing error and the args, result and value are not correctly infered when possible. See the playground attached.

Reproduction Link

https://www.typescriptlang.org/play/?#code/JYWwDg9gTgLgBAbzgQxhEAaOBjKBTVPAZTXzgF84AzKdOAIgCsIZlh6AoD7CAOwGd4g6HjgBeHPkIkRACgCUXHgPjJxKNCFkBGRcsFwARutTpZAckPnF3PgeQBVYxICyqABYA6KMl4ATM3k4AD44AAZPAFY4AH4UOAAuIzgAehSNdAAeTxzQgB8MkGzcpTt4AFcANWQAG3LRCWF8TwBzPBhZR0MgtLgYAE8wUTwoWigUfzgq2vq4YH44XhY53ioRvD8UBd5ykEMRuALBKGBeFq4mvE9+ds6nLAVxUPNePAB3a1T0gaG4EbGuL1ru0UFAWgtfJt8DBylBeHAAG4zUSnNZQUoqHAmTSZHZ7EZYADaeP2UAAulgSSNgjoHm0YFgbgzFkExKFdLZMZsJKYilSoETjqcWhS4EKzjTtHT2oyZSynnBzAFXtZOQZsA5uXA3DAvD5-IEQuEorEsUk-BdSFcmbINX4sOZkJ9ej9hqNoBxLsCOnasLovn1Bm6Afp4ONGlbvbbNQ6nT10uN5is0RstotdqSJptgEJ3BByjVNvt0-jxkcYCczhwgA

[double-thinker]

This commit updates Updated PrdStore to use the ExtractAtomValue, ExtractAtomArgs and ExtractAtomResult utility types. It fixes the inferring and removes the errors.

[dai-shi]

Yeah, it's not supported at the moment. We could do something similar to useAtomValue type.

[double-thinker]

Are you willing to add them if I create a PR with the same changes for ùseAtomValue?

[dai-shi]

I think main...double-thinker:jotai:bugfix/atom-unions looks okay.
Do you know if there's any drawback with the change?

[double-thinker]

I am not aware of any drawback, but not a TypeScript expert. I ran all the tests, and tsc does not throw any errors. To the best of my understanding, making it more generic should not break anything backward-compatible, and in the worst-case scenario, it would not affect runtime behavior at all.

I will make changes to useAtomValues and submit a pull request.