Code tidying up