Improve Python script