Made some small refactorings
Merge pull request #2 from StevenMcD/minor_refactor

Thanks @StevenMcD