Tweaked code for rubocop