Minor Refactor

Tidied up code