Updated Python code