Hope it closes #2
Minor refactoring