Add callbacks and image task to WKInterfaceImage extension

Sources/WKInterfaceImage+Kingfisher.swift

@@ -35,17 +35,87 @@ extension Kingfisher where Base: WKInterfaceImage {
      Set an image with a resource.
 
      - parameter resource:          Resource object contains information such as `cacheKey` and `downloadURL`.
+     - parameter options:           A dictionary could control some behaviors. See `KingfisherOptionsInfo` for more.
+     - parameter progressBlock:     Called when the image downloading progress gets updated.
+     - parameter completionHandler: Called when the image retrieved and set.
+
+     - returns: A task represents the retrieving process.
      */
-    public func setImage(_ resource: Resource?) {
+    @discardableResult
+    public func setImage(_ resource: Resource?,
+                         options: KingfisherOptionsInfo? = nil,
+                         progressBlock: DownloadProgressBlock? = nil,
+                         completionHandler: CompletionHandler? = nil) -> RetrieveImageTask {
         guard let resource = resource else {
-            return
+            setWebURL(nil)
+            completionHandler?(nil, nil, .none, nil)
+            return .empty
         }
 
-        KingfisherManager.shared.retrieveImage(with: resource, options: nil, progressBlock: nil, completionHandler: { [weak base] image, error, cacheType, imageURL in
-            DispatchQueue.main.safeAsync {
-                base?.setImage(image)
-            }
-        })
+        setWebURL(resource.downloadURL)
+
+        let task = KingfisherManager.shared.retrieveImage(
+            with: resource,
+            options: KingfisherManager.shared.defaultOptions + (options ?? KingfisherEmptyOptionsInfo),
+            progressBlock: { receivedSize, totalSize in
+                guard resource.downloadURL == self.webURL else {
+                    return
+                }
+
+                progressBlock?(receivedSize, totalSize)
+            },
+            completionHandler: { [weak base] image, error, cacheType, imageURL in
+                DispatchQueue.main.safeAsync {
+                    guard let strongBase = base, imageURL == self.webURL else {
+                        completionHandler?(image, error, cacheType, imageURL)
+                        return
+                    }
+
+                    self.setImageTask(nil)
+                    guard let image = image else {
+                        completionHandler?(nil, error, cacheType, imageURL)
+                        return
+                    }
+                
+                    strongBase.setImage(image)
+                    completionHandler?(image, error, cacheType, imageURL)
+                }
+            })
+
+        setImageTask(task)
+
+        return task
+    }
+
+    /**
+     Cancel the image download task bounded to the image view if it is running.
+     Nothing will happen if the downloading has already finished.
+     */
+    public func cancelDownloadTask() {
+        imageTask?.cancel()
+    }
+}
+
+// MARK: - Associated Object
+private var lastURLKey: Void?
+private var imageTaskKey: Void?
+
+extension Kingfisher where Base: WKInterfaceImage {
+    /// Get the image URL binded to this image view.
+    public var webURL: URL? {
+        return objc_getAssociatedObject(base, &lastURLKey) as? URL
+    }
+
+    fileprivate func setWebURL(_ url: URL?) {
+        objc_setAssociatedObject(base, &lastURLKey, url, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
+    }
+
+    fileprivate var imageTask: RetrieveImageTask? {
+        return objc_getAssociatedObject(base, &imageTaskKey) as? RetrieveImageTask
+    }
+
+    fileprivate func setImageTask(_ task: RetrieveImageTask?) {
+        objc_setAssociatedObject(base, &imageTaskKey, task, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
     }
 }