* Improved type hints for when() helper function.

* Simplify conditional return type of `when` helper function

Co-authored-by: Sebastian Hädrich <11225821+shaedrich@users.noreply.github.com>

---------

Co-authored-by: Sebastian Hädrich <11225821+shaedrich@users.noreply.github.com>