A bit of a tidy up

No tests, so might not work! It should though @pikesley, could you take a look?