Improve type hints

Change-Id: I05ba765ca721877292bf2e3125ec8e42a95876ac
Merge "Improve type hints"