Fix missing import, pep8