Refactored code to DRY it up a bit