Add missing import

Found thanks to flake8