Code tidy ups