FirebaseStorage with RxSwift

RxFireStorage.swift

//
//  RxFireStorage.swift
//  RxSwiftSamples
//
//  Created by Yuhei Miyazato on 2017/11/25.
//  Copyright © 2017 mitolab. All rights reserved.
//

import Foundation
import FirebaseStorage
import RxSwift

// customized version of
// https://github.com/RxSwiftCommunity/RxFirebase/blob/master/RxFirebase/Classes/RxFirebaseStorage.swift

extension Reactive where Base : Storage {

    // must be single to make observable .complete state, otherwise it leaks
    public static func putData(data: Data, to path:String, metaData: StorageMetadata? = nil) -> Single<StorageUploadTask> {
        return Single.create { observer in
            observer(.success(Base.storage().reference().child(path).putData(data, metadata: metaData)))
            return Disposables.create()
        }
    }

    public static func putDataWithProgress(data: Data, to path:String, metaData: StorageMetadata? = nil) -> Observable<(StorageTaskSnapshot, StorageTaskStatus)> {
        return self.putData(data: data, to:path, metaData: metaData).asObservable().storageObserve()
    }
}

extension Reactive where Base : StorageUploadTask {
    
    func observeStatus(_ status: StorageTaskStatus) -> Observable<(StorageTaskSnapshot, StorageTaskStatus)> {
        return Observable.create { observer in
            let observeStatus = self.base.observe(status, handler: { snapshot in
                
                if let error = snapshot.error {
                    observer.onError(error)
                } else {
                    // There might be a case (status == .failure && snapshot.error == nil)? don't know...
                    // This have to be checked in the future.
                    observer.onNext((snapshot, status))
                    if status == .success {
                        observer.onCompleted()
                    }
                }
            })
            return Disposables.create {
                self.base.removeObserver(withHandle: observeStatus)
            }
        }
    }
}

extension ObservableType where E : StorageUploadTask {
    
    func storageObserve() -> Observable<(StorageTaskSnapshot, StorageTaskStatus)> {
        return self.flatMap { (uploadTask: E) -> Observable<(StorageTaskSnapshot, StorageTaskStatus)> in
            
            // bit messy... but it's necessary to prevent leaks.
            let success = uploadTask.rx.observeStatus(.success)
            let fail = uploadTask.rx.observeStatus(.failure)
            let successStatus = success.takeUntil(fail)
            let failureStatus = fail.takeUntil(success)
            let progressStatus = uploadTask.rx.observeStatus(.progress).takeWhile({ (snapshot, _) -> Bool in
                !snapshot.progress!.isFinished
            })
            
            return Observable.of(progressStatus, successStatus, failureStatus).merge()
        }
    }
}