Clean up python script