Python if fix