Add better coloring for comments