Making NotificationCenter easier to work with and allowing type-safe rules

NotificationCenter.swift

extension NotificationCenter {

    public func observe(
        name: NSNotification.Name,
        object obj: Any? = nil,
        queue: OperationQueue? = .main,
        using block: @escaping (Notification?) -> Void
    ) -> Notification.Token {
        return Notification.Token(
            on: self,
            token: addObserver(forName: name, object: obj, queue: queue, using: block)
        )
    }
}

extension NotificationCenter {
    
    public func observe<T>(
        _ descriptor: Notification.Descriptor<T>,
        object obj: Any? = nil,
        queue: OperationQueue? = .main,
        using block: @escaping (T) -> ()
    ) -> Notification.Token {
        return Notification.Token(
            on: self,
            token: addObserver(for: descriptor, object: obj, queue: queue, using: block)
        )
    }
    
    public func addObserver<T>(
        for descriptor: Notification.Descriptor<T>,
        object obj: Any? = nil,
        queue: OperationQueue? = .main,
        using block: @escaping (T) -> ()
    ) -> NSObjectProtocol {
        return addObserver(forName: descriptor.name, object: obj, queue: queue, using: { n in
            block(descriptor.transform(n))
        })
    }
}

extension Notification {
    
    public final class Token: NSObject {
        
        let notificationCenter: NotificationCenter
        let token: NSObjectProtocol
        
        init(on notificationCenter: NotificationCenter = .default, token: NSObjectProtocol) {
            self.notificationCenter = notificationCenter
            self.token = token
        }
        
        deinit {
            notificationCenter.removeObserver(token)
        }
    }
    
    public struct Descriptor<T> {
        
        public typealias ID = UInt
        private(set) public var id: ID = { __count &+= 1; return __count }()
        
        public let name: Notification.Name
        public let transform: (Notification) -> T
        
        public init(name: Notification.Name, transform: @escaping (Notification) -> T) {
            self.name = name
            self.transform = transform
        }
    }
}

private var __count: UInt = 0

extension Notification.Descriptor: Hashable {
    
    public func hash(into hasher: inout Hasher) {
        hasher.combine(name)
    }
    
    public static func == (l: Notification.Descriptor<T>, r: Notification.Descriptor<T>) -> Bool {
        l.name == r.name && l.id == r.id
    }
}

extension Notification.Descriptor {
    
    public func sink<Root>(center: NotificationCenter = .default, queue: OperationQueue? = nil, to method: @escaping (Root) -> (T) -> Void, on root: Root) -> Notification.Token where Root: AnyObject {
        sink(center: center, queue: queue, to: { [weak root] o in root.map(method)?(o) })
    }
    public func sink(center: NotificationCenter = .default, queue: OperationQueue? = nil, to method: @escaping (T) -> Void) -> Notification.Token {
        center.observe(self, object: nil, queue: queue, using: method)
    }
}

extension Notification.Name {
    
    public func sink<Root>(center: NotificationCenter = .default, queue: OperationQueue? = nil, to method: @escaping (Root) -> () -> Void, on root: Root) -> Notification.Token where Root: AnyObject {
        sink(center: center, queue: queue, to: { [weak root] _ in root.map(method)?() })
    }
    public func sink(center: NotificationCenter = .default, queue: OperationQueue? = nil, to method: @escaping () -> Void) -> Notification.Token {
        center.observe(name: self, object: nil, queue: queue, using: { _ in method() })
    }
    public func sink<Root>(center: NotificationCenter = .default, queue: OperationQueue? = nil, to method: @escaping (Root) -> (Notification?) -> Void, on root: Root) -> Notification.Token where Root: AnyObject {
        sink(center: center, queue: queue, to: { [weak root] note in root.map(method)?(note) })
    }
    public func sink(center: NotificationCenter = .default, queue: OperationQueue? = nil, to method: @escaping (Notification?) -> Void) -> Notification.Token {
        center.observe(name: self, object: nil, queue: queue, using: method)
    }
}

// AnyToken

public final class AnyObservationToken: NSObject {
    let wrapped: NSObject
    public init(_ wrapped: NSObject) {
        self.wrapped = wrapped
    }
}

protocol ExpressibleAsAnyObservationToken: NSObject {
    func store(in set: inout Set<AnyObservationToken>)
    func store<C>(in collection: inout C) where C: RangeReplaceableCollection, C.Element == AnyObservationToken
    func eraseToAnyObservationToken() -> AnyObservationToken
}

extension ExpressibleAsAnyObservationToken {
    public func store(in set: inout Set<AnyObservationToken>) {
        set.insert(eraseToAnyObservationToken())
    }
    public func store<C>(in collection: inout C) where C : RangeReplaceableCollection, C.Element == AnyObservationToken {
        collection.append(eraseToAnyObservationToken())
    }
    public func eraseToAnyObservationToken() -> AnyObservationToken { .init(self) }
}

extension AnyObservationToken {
    public func store(in set: inout Set<AnyObservationToken>) {
        set.insert(self)
    }
}

extension NSKeyValueObservation: ExpressibleAsAnyObservationToken { }
extension Notification.Token: ExpressibleAsAnyObservationToken { }

// KeyPath + Unwrap

extension KeyPath {
    
    public func tryUnwrap<T>(_ as: T.Type = T.self) -> (Root) throws -> T {
        return { try $0[keyPath: self] as? T ??^ "\(self) is not type \(T.self)" }
    }
    
    public func optionalUnwrap<T>(_ as: T.Type = T.self) -> (Root) -> T? {
        return { $0[keyPath: self] as? T }
    }
    
    public func unwrap<T>(_ as: T.Type = T.self) -> (Root) -> T {
        return { $0[keyPath: self] as! T }
    }
}

// Unwrap or throw

extension String: Error { }

extension Optional {
    public static func ??^ (optional: Optional, error: Error) throws -> Wrapped {
        guard case let value? = optional else { throw error }
        return value
    }
}

UIWindow+Notification.swift

extension UIWindow {

    public static let didBecomeVisibleNotificationDescriptor: Notification.Descriptor<UIWindow> = .init(name: UIWindow.didBecomeVisibleNotification, transform: (\Notification.object).unwrap())
    public static let didBecomeHiddenNotificationDescriptor: Notification.Descriptor<UIWindow> = .init(name: UIWindow.didBecomeHiddenNotification, transform: (\Notification.object).unwrap())

    public static let didBecomeKeyNotificationDescriptor: Notification.Descriptor<UIWindow> = .init(name: UIWindow.didBecomeKeyNotification, transform: (\Notification.object).unwrap())
    public static let didResignKeyNotificationDescriptor: Notification.Descriptor<UIWindow> = .init(name: UIWindow.didResignKeyNotification, transform: (\Notification.object).unwrap())
}

e.g.

private var bag: Set<AnyToken> = [ ]

...

UIWindow.didBecomeKeyNotificationDescriptor
    .sink(to: windowDidBecomeKey(_:))
    .store(in: &bag)

...

func windowDidBecomeKey(_ window: UIWindow) {
    ...
}

or grabbing properties from the notification

extension UIResponder {
    static let keyboardDidChangeFrameNotificationDescriptor = Notification.Descriptor(name: UIResponder.keyboardDidChangeFrameNotification, transform: KeyboardDidChangeFrame.init)
}

struct KeyboardDidChangeFrame {
    let from: CGRect
    let to: CGRect
    init(_ note: Notification) {
        from = note.userInfo?[UIResponder.keyboardFrameBeginUserInfoKey] as! CGRect
        to = note.userInfo?[UIResponder.keyboardFrameEndUserInfoKey] as! CGRect
    }
}

...

UIResponser.keyboardDidChangeFrameNotificationDescriptor
    .sink(to: keyboardDidChangeFrame(_:))
    .store(in: &bag)

...

func keyboardDidChangeFrame(_ change: KeyboardDidChangeFrame) {
    ...
}