Minor change to improve code from #122

This is the first time I've tried to push back to a PR, which isn't possible (makes sense, it's not my repo).

/cc @scottlangendyk