A badge label for iOS

BadgeLabel.swift

import UIKit

/**
 * [Badge.swift](https://gist.github.com/yonat/75a0f432d791165b1fd6) をベース
 */
class BadgeLabel: UILabel {
    var badgeColor: UIColor = UIColor.redColor()

    convenience init(badge: String?, color: UIColor = UIColor.redColor()) {
        self.init()
        badgeColor = color
        text = badge
    }

    override init(frame: CGRect) {
        super.init(frame: frame)
        initialize()
    }

    required init?(coder aDecoder: NSCoder) {
        super.init(coder: aDecoder)
        initialize()
    }

    func initialize() {
        textColor = UIColor.whiteColor()
        backgroundColor = UIColor.clearColor()
        textAlignment = .Center
    }

    override func drawRect(rect: CGRect) {
        let context = UIGraphicsGetCurrentContext()
        CGContextSetFillColorWithColor(context, badgeColor.CGColor)
        CGContextFillEllipseInRect(context, rect)
        super.drawRect(rect)
    }

    override var text: String? {
        didSet {
            guard let text = self.text else {
                hidden = true
                return
            }
            if text.isEmpty {
                hidden = true
                return
            }
            if let count = Int(text) where count == 0 {
                hidden = true
                return
            }
            hidden = false
        }
    }

    let margin: CGFloat = 3

    override func drawTextInRect(rect: CGRect) {
        let insets = UIEdgeInsets(top: margin, left: margin, bottom: margin, right: margin)
        let rect = UIEdgeInsetsInsetRect(rect, insets)
        super.drawTextInRect(rect)
    }

    override func intrinsicContentSize() -> CGSize {
        var size = super.intrinsicContentSize()
        let width = size.width + margin * 2
        let height = size.height + margin * 2
        if width >= height {
            size.width = width
            size.height = width
        } else {
            size.width = height
            size.height = height
        }
        return size
    }
}

extension UIBarButtonItem {
    convenience init(image: UIImage?, badge: String?, target: AnyObject?, action: Selector) {
        let button = UIButton(type: .Custom)
        button.frame = CGRect(origin: CGPointZero, size: image?.size ?? CGSizeZero)
        button.setImage(image, forState: .Normal)
        button.addTarget(target, action: action, forControlEvents: .TouchUpInside)
        self.init(customView: button)

        let label = BadgeLabel(badge: badge)
        label.tag = self.dynamicType.badgeTag
        button.addSubview(label)

        label.translatesAutoresizingMaskIntoConstraints = false
        button.addConstraint(NSLayoutConstraint(item: label, attribute: .CenterX, relatedBy: .Equal, toItem: button, attribute: .CenterX, multiplier: 1, constant: 10))
        button.addConstraint(NSLayoutConstraint(item: button, attribute: .CenterY, relatedBy: .Equal, toItem: label, attribute: .CenterY, multiplier: 1, constant: 10))
    }

    var badgeButton: UIButton {
        return self.customView as! UIButton
    }

    var badgeLabel: BadgeLabel {
        return self.customView?.viewWithTag(self.dynamicType.badgeTag) as! BadgeLabel
    }

    var badgeText: String? {
        get { return badgeLabel.text }
        set {
            badgeLabel.text = newValue
            badgeLabel.layoutIfNeeded()
        }
    }

    private static let badgeTag = 8181
}