[UPDATE] Pythonic code