Make code more compact

I think this is better; before you had two blocks of code that did exactly the same
Merge pull request #9 from Sunnywout/patch-1

Make code more compact