Cleaning code - part 2