[UIBarButtonItem+Badge] #iOS #extension

UIBarButtonItem+Badge.swift

//
//  UIBarButtonItem+Badge.swift
//
//  Created by Jaeho Lee on 30/08/2019.
//  Copyright © 2019 Jay Lee <jaeho.lee@snu.ac.kr>
//
//  Original Source:
//  https://gist.github.com/freedom27/c709923b163e26405f62b799437243f4
//

import UIKit

extension CAShapeLayer {

  func drawRoundedRect(rect: CGRect, andColor color: UIColor, filled: Bool) {
    fillColor = filled ? color.cgColor : UIColor.white.cgColor
    strokeColor = color.cgColor
    path = UIBezierPath(roundedRect: rect, cornerRadius: 8.5).cgPath
  }
}

private var handle: UInt8 = 0

extension UIBarButtonItem {

  private var badgeLayer: CAShapeLayer? {
    if let b = objc_getAssociatedObject(self, &handle) as AnyObject? {
      return b as? CAShapeLayer
    } else {
      return nil
    }
  }


  func setBadge(
    text: String?,
    offset: CGPoint = .zero,
    color: UIColor = .red,
    filled: Bool = true,
    font: UIFont = .systemFont(ofSize: UIFont.systemFontSize))
  {
    badgeLayer?.removeFromSuperlayer()
    guard let text = text, !text.isEmpty else { return }
    addBadge(
      text: text,
      offset: offset,
      color: color,
      filled: filled,
      font: font)
  }


  private func addBadge(
    text: String,
    offset: CGPoint,
    color: UIColor,
    filled: Bool,
    font: UIFont)
  {
    guard let view = self.value(forKey: "view") as? UIView else { return }

    let badgeSize = text.size(withAttributes: [.font: font])

    // Initialize Badge
    let badge = CAShapeLayer()

    let height = badgeSize.height;
    var width = badgeSize.width + 2 /* padding */

    // Make sure we have at least a circle
    if (width < height) { width = height }

    // x position is offset from right-hand side
    let x = view.frame.width - width + offset.x

    let badgeFrame = CGRect(origin: CGPoint(x: x, y: offset.y), size: CGSize(width: width, height: height))

    badge.drawRoundedRect(rect: badgeFrame, andColor: color, filled: filled)
    view.layer.addSublayer(badge)

    // Initialize the label of the badge
    let label = CATextLayer()
    label.string = text
    label.alignmentMode = .center
    label.font = font
    label.fontSize = font.pointSize

    label.frame = badgeFrame
    label.foregroundColor = filled ? UIColor.white.cgColor : color.cgColor
    label.backgroundColor = UIColor.clear.cgColor
    label.contentsScale = UIScreen.main.scale
    badge.addSublayer(label)

    badge.zPosition = CGFloat(Float.greatestFiniteMagnitude)

    // Save badge as a UIBarButtonItem property
    objc_setAssociatedObject(self, &handle, badge, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
  }
}