Fix script for python 2.7