Fix reverted pull request 2

@kobynet

Can you please fix second pull request on this branch, with regards to my comments on #2?
I had to open a new pull request since you merged #2 but it wasn't ready so I reverted it.