Updates to code from pull request comments