diff --git a/TDBadgedCell/TDBadgedCell.swift b/TDBadgedCell/TDBadgedCell.swift
index 7a755c0..a2e6db2 100644
--- a/TDBadgedCell/TDBadgedCell.swift
+++ b/TDBadgedCell/TDBadgedCell.swift
@@ -43,6 +43,12 @@ class TDBadgedCell: UITableViewCell {
         
         badgeView.frame.origin.x = floor(self.contentView.frame.width - badgeView.frame.width - badgeOffset.x)
         badgeView.frame.origin.y = floor((self.frame.height / 2) - (badgeView.frame.height / 2))
+        
+        // Now lets update the width of the cells text labels to take the badge into account
+        self.textLabel?.frame.size.width -= badgeView.frame.width + (badgeOffset.x * 2)
+        if((self.detailTextLabel) != nil) {
+            self.detailTextLabel?.frame.size.width -= badgeView.frame.width + (badgeOffset.x * 2)
+        }
     }
     
     // When the badge
@@ -92,7 +98,7 @@ class TDBadgedCell: UITableViewCell {
         NSString(string: badgeString).draw(in:CGRect(x:8, y:5, width:textSize.width, height:textSize.height), withAttributes: [
             NSFontAttributeName:UIFont.boldSystemFont(ofSize:CGFloat(badgeFontSize)),
             NSForegroundColorAttributeName: UIColor.clear
-            ])
+        ])
         
         let badgeImage = UIGraphicsGetImageFromCurrentImageContext()!
         UIGraphicsEndImageContext()
diff --git a/TDBadgedCell/ViewController.swift b/TDBadgedCell/ViewController.swift
index 0a96c08..bee1591 100644
--- a/TDBadgedCell/ViewController.swift
+++ b/TDBadgedCell/ViewController.swift
@@ -14,23 +14,14 @@ class ViewController: UITableViewController {
         ["title" : "This is an example badge", "badge": "1"],
         ["title" : "This is a second example badge", "badge": "123"],
         ["title" : "A text badge", "badge": "Warning!"],
-        ["title" : "Another text badge", "badge": "Danger!"],
+        ["title" : "Another text badge with a really long title!", "badge": "Danger!"],
     ]
 
-    override func viewDidLoad() {
-        super.viewDidLoad()
-        self.tableView.delegate = self
-        self.tableView.dataSource = self
-    }
     
     override func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
         return demoItems.count
     }
     
-    override func numberOfSections(in tableView: UITableView) -> Int {
-        return 1
-    }
-    
     override func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> TDBadgedCell {
         var cell : TDBadgedCell? = tableView.dequeueReusableCell(withIdentifier:"BadgedCell") as! TDBadgedCell?;
         if((cell == nil)) {