Improve output highlighting