Small code refactor