Fixed code style, passes flake8