improved python script