Updated python code