UITextViewへの入力時の文字属性の変更について

UITextViewに下記のようなRichText文字列が表示されているとします。
ああ　　　いい　　　うう
↑黒色　　↑赤色　　　↑青色

たとえば、緑色指定ボタンを事前に押して、「うう」の前に、
緑色でテキストを入力したい場合は、
どうすれば入力時の文字色の指定ができるのでしょうか？

一部文字色の変更は、下記のような実装で実現できますが、
入力時の色指定方法が分かりません。

var range: NSRange = NSRange(location: 1, length: 1)
var attributedString =     NSMutableAttributedString(attributedString:self.testUiText.attributedText)
attributedString.addAttribute(NSForegroundColorAttributeName, value: UIColor.redColor(), range: range)
self.testUiText.attributedText = attributedString

↑self.testUiText.attributedTextに代入すると、キャレットが一番後ろに
なってしまいます。

func textView(textView: UITextView, shouldChangeTextInRange range: NSRange, replacementText text: String) -> Bool

や

func textViewDidChange(textView: UITextView)

で処理できるのでしょうか？