Made code shorter using some Python idioms