[mypy] Add type hints and docstrings to heap.py (#3013)

* Add type hints and docstrings to heap.py

- Add type hints
- Add docstrings
- Add explanatory comments 
- Improve code readability
- Change to use f-string

* Fix import sorting

* fixup! Format Python code with psf/black push

* Fix static type error

* Fix failing test

* Fix type hints

* Add return annotation

Co-authored-by: github-actions <${GITHUB_ACTOR}@users.noreply.github.com>
Co-authored-by: Dhruv Manilawala <dhruvmanila@gmail.com>