Merge pull request #36 from GlisboaDev/dismiss_handler

Added handler for dismiss view action.
itsmeichigo · Jul 28, 2017 · 3c122d8 · 3c122d8
2 parents e2bc255 + 7d61744
commit 3c122d8
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/Source/DateTimePicker.swift b/Source/DateTimePicker.swift
@@ -101,6 +101,7 @@ import UIKit
 
     public var timeZone = TimeZone.current
     public var completionHandler: ((Date)->Void)?
+    public var dismissHandler: (() -> Void)?
 
     // private vars
     internal var hourTableView: UITableView!
@@ -459,9 +460,14 @@ import UIKit
                                             y: self.frame.height,
                                             width: self.frame.width,
                                             height: self.contentHeight)
-        }) { (completed) in
+        }) {[weak self] (completed) in
+            guard let `self` = self else {
+                return
+            }
             if sender == self.doneButton {
                 self.completionHandler?(self.selectedDate)
+            } else {
+                self.dismissHandler?()
             }
             self.removeFromSuperview()
         }