Fix return type hints, and cleanup imports.