Fixed issues reported in #4 in addition DRY'ed a little bit