Add suggested changes from code review in #338