Refactor to be more Python-y, as far as I understand =P