Refactor in response to pull request comments