Tidying up code some more