Remove ad-hoc code and flake8 cleanup