cleanup - results from autopep8

Just in the main module