Refactor/clean up of python code.