CoreData NSManagedObjectContext observer

CoreDataContextObserver.swift

import Foundation
import CoreData

public class CoreDataContextObserver {
    public typealias ActionCompletionBlock = (NSManagedObject, State) -> Void
    public typealias ContextChangeBlock = (Foundation.Notification, [ObjectChange]) -> Void

    public struct State: OptionSet {
        public let rawValue: Int

        public init(rawValue: Int) {
            self.rawValue = rawValue
        }

        public static let inserted = State(rawValue: 1 << 0)
        public static let updated = State(rawValue: 1 << 1)
        public static let deleted = State(rawValue: 1 << 2)
        public static let refreshed = State(rawValue: 1 << 3)
        public static let all: State = [inserted, updated, deleted, refreshed]
    }

    public struct Action {
        var state: State
        var completionBlock: ActionCompletionBlock
    }

    public enum ObjectChange {
        case updated(NSManagedObject)
        case refreshed(NSManagedObject)
        case inserted(NSManagedObject)
        case deleted(NSManagedObject)

        public func managedObject() -> NSManagedObject {
            switch self {
            case let .updated(value):
                return value
            case let .inserted(value):
                return value
            case let .refreshed(value):
                return value
            case let .deleted(value):
                return value
            }
        }
    }

    public var enabled: Bool = true
    public var contextChangeBlock: ContextChangeBlock?
    private var notificationObserver: NSObjectProtocol?

    private(set) var actionsMap: Dictionary<NSManagedObjectID, [Action]> = [:]
    private(set) var context: NSManagedObjectContext
    private(set) weak var persistentStoreCoordinator: NSPersistentStoreCoordinator?

    public init(context: NSManagedObjectContext) {
        self.context = context
        self.persistentStoreCoordinator = context.persistentStoreCoordinator

        notificationObserver = NotificationCenter.default.addObserver(
            forName: .NSManagedObjectContextObjectsDidChange,
            object: context,
            queue: nil) { [weak self] notification in
            self?.handleNotification(notification)
        }
    }

    deinit {
        removeAll()
        if let notificationObserver = notificationObserver {
            NotificationCenter.default.removeObserver(notificationObserver)
        }
    }

    // MARK: - Public API -

    public func observeObject(_ object: NSManagedObject, state: State = .all, completionBlock: @escaping ActionCompletionBlock) {
        let action = Action(state: state, completionBlock: completionBlock)
        if var actionArray = actionsMap[object.objectID] {
            actionArray.append(action)
            actionsMap[object.objectID] = actionArray
        } else {
            actionsMap[object.objectID] = [action]
        }
    }

    public func removeObserver(for object: NSManagedObject, withState state: State = .all) {
        if state == .all {
            actionsMap.removeValue(forKey: object.objectID)
        } else if let actionsForObject = actionsMap[object.objectID] {
            actionsMap[object.objectID] = actionsForObject.filter({ !$0.state.contains(state) })
        }
    }

    public func removeAll() {
        actionsMap.removeAll()
    }

    // MARK: - Private API -

    private func handleNotification(_ notification: Foundation.Notification) {
        guard let incomingContext = notification.object as? NSManagedObjectContext,
              let persistentStoreCoordinator = persistentStoreCoordinator,
              let incomingPersistentStoreCoordinator = incomingContext.persistentStoreCoordinator,
              enabled && persistentStoreCoordinator == incomingPersistentStoreCoordinator else {
            return
        }

        let insertedObjectsSet = notification.userInfo?[NSInsertedObjectsKey] as? Set<NSManagedObject> ?? Set<NSManagedObject>()
        let updatedObjectsSet = notification.userInfo?[NSUpdatedObjectsKey] as? Set<NSManagedObject> ?? Set<NSManagedObject>()
        let deletedObjectsSet = notification.userInfo?[NSDeletedObjectsKey] as? Set<NSManagedObject> ?? Set<NSManagedObject>()
        let refreshedObjectsSet = notification.userInfo?[NSRefreshedObjectsKey] as? Set<NSManagedObject> ?? Set<NSManagedObject>()

        var combinedObjectChanges = insertedObjectsSet.map({ ObjectChange.inserted($0) })
        combinedObjectChanges += updatedObjectsSet.map({ ObjectChange.updated($0) })
        combinedObjectChanges += deletedObjectsSet.map({ ObjectChange.deleted($0) })
        combinedObjectChanges += refreshedObjectsSet.map({ ObjectChange.refreshed($0) })

        contextChangeBlock?(notification, combinedObjectChanges)

        let combinedSet = insertedObjectsSet.union(updatedObjectsSet).union(deletedObjectsSet)
        let allObjectIDs = Array(actionsMap.keys)
        let filteredObjects = combinedSet.filter({ allObjectIDs.contains($0.objectID) })

        for object in filteredObjects {
            guard let actionsForObject = actionsMap[object.objectID] else {
                continue
            }

            for action in actionsForObject {
                if action.state.contains(.inserted) && insertedObjectsSet.contains(object) {
                    action.completionBlock(object, .inserted)
                } else if action.state.contains(.updated) && updatedObjectsSet.contains(object) {
                    action.completionBlock(object, .updated)
                } else if action.state.contains(.deleted) && deletedObjectsSet.contains(object) {
                    action.completionBlock(object, .deleted)
                } else if action.state.contains(.refreshed) && refreshedObjectsSet.contains(object) {
                    action.completionBlock(object, .refreshed)
                }
            }
        }
    }

}