Better colon handling

Previous version of the script was not handling cases like

if x == 2: # Some comment
<newline here>

properly. Now, it does.