Improve docs and type hints (#5759)
* improvement: docs and type hints
* improvement: docs and type hints
minor refactor
* improvement: docs and type hints
* update with suggestions from review
Co-Authored-By: Dhruv Nair <dhruv.nair@gmail.com>
---------
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
