py script improvement