code refactoring for #5 - as @alanpearce solution is nice