[NSObject+AssociatedObject] Wrapper around associated objects #swift

NSObject+AssociatedObject.swift

public enum AssociationPolicy {
    case assign
    case retain
    case copy
    case retainNonatomic
    case copyNonatomic
    
    fileprivate var policy: objc_AssociationPolicy {
        switch self {
            case .assign: return .OBJC_ASSOCIATION_ASSIGN
            case .retain: return .OBJC_ASSOCIATION_RETAIN
            case .copy: return .OBJC_ASSOCIATION_COPY
            case .retainNonatomic: return .OBJC_ASSOCIATION_RETAIN_NONATOMIC
            case .copyNonatomic: return .OBJC_ASSOCIATION_COPY_NONATOMIC
        }
    }
}

public extension NSObjectProtocol {
    
    public func setAssociatedObject(_ object: Any?, forKey key: UnsafeRawPointer, policy: AssociationPolicy = .retain) {
        objc_setAssociatedObject(self, key, object, policy.policy)
    }
    
    public func associatedObject<T>(for key: UnsafeRawPointer) -> T? {
        return objc_getAssociatedObject(self, key) as? T
    }
    
}