Python script improvement