add init for UIBarButtonItem

jogendra · Oct 19, 2019 · 5549538 · 5549538
1 parent 3c70f6a
commit 5549538
Showing 1 changed file with 13 additions and 7 deletions.
diff --git a/BadgeHub/Classes/BadgeHub.swift b/BadgeHub/Classes/BadgeHub.swift
@@ -67,13 +67,19 @@ public class BadgeHub: NSObject {
         setView(view, andCount: 0)
     }
 
-    //    convenience init(barButtonItem: UIBarButtonItem?) {
-    //        if let value = barButtonItem?.value(forKey: "view") as? UIView {
-    //            self.init(view: value)
-    //            scaleCircleSize(by: 0.7)
-    //            moveCircleBy(x: -5.0, y: 0)
-    //        }
-    //    }
+    public convenience init?(barButtonItem: UIBarButtonItem) {
+        if let value = barButtonItem.value(forKey: "view") as? UIView {
+            self.init(view: value)
+            scaleCircleSize(by: 0.7)
+            moveCircleBy(x: -5.0, y: 0)
+        } else if let value = barButtonItem.customView {
+            self.init(view: value)
+            scaleCircleSize(by: 0.7)
+            moveCircleBy(x: -5.0, y: 0)
+        } else {
+            return nil
+        }
+    }
 
     // Adjustment methods
     public func setView(_ view: UIView?, andCount startCount: Int) {