Small code tidy