Finishing what was started with #65

I applied the fix that #65 adds and applied the feedback in the comments.