UIButton subclass with per-state custom values like background and image colors, extendable with whatever value you want to update

CustomButton.swift

class Button: UIButton {
    
    private class Subview: UIView {
        // Never allow isUserInteractionEnabled in any button subview
        override var isUserInteractionEnabled: Bool {
            get { false }
            set { super.isUserInteractionEnabled = false }
        }
    }
    
    // Creating custom classes for each subview for inspection purposes
    private class BackgroundView: Subview { }
    private class ShadowView: Subview { }
    
    var animatesStateChanges: Bool = false
    
    private(set) lazy var backgroundView: UIView = BackgroundView(frame: bounds)
    private(set) lazy var shadowView: UIView = ShadowView(frame: bounds)
    
    private var originalTintColor: UIColor!
    
    private var stateValueContainers: [PartialKeyPath<Button> : AnyStateValueContainer] = [:]
    
    override init(frame: CGRect) {
        super.init(frame: frame)
        originalTintColor = tintColor
        insertSubview(backgroundView, at: 0)
        insertSubview(shadowView, belowSubview: backgroundView)
    }
    
    required init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
}

// MARK: - Updating apperance
extension Button {
    
    override var isHighlighted: Bool { didSet {
        updateAppearance(animated: animatesStateChanges)
    }}
    
    override var isEnabled: Bool { didSet {
        updateAppearance()
    }}
    
    override var isSelected: Bool { didSet {
        updateAppearance(animated: animatesStateChanges)
    }}
    
    func updateAppearance(animated: Bool = false) {
        
        let updates = {
            // All self-animating value containers or all updates when non animated
            self.stateValueContainers
                .filter { !$1.animatesUpdate || !animated }
                .forEach { keyPath, container in
                    container.update(object: self, for: self.state, using: keyPath)
                }
        }
        if animated {
            UIView.animate(
                withDuration: 0.25, delay: 0,
                options: [.beginFromCurrentState, .allowUserInteraction],
                timingFunction: .quintOut,
                animations: updates)
            
            // Let all self-animating value containers animate
            self.stateValueContainers
                .filter { $1.animatesUpdate }
                .forEach { keyPath, container in
                    container.update(object: self, for: self.state, using: keyPath)
                }
        } else {
            UIView.performWithoutAnimation {
                // Disable all self-animating value container animations
                updates()
            }
        }
        
    }
}

// MARK: - View hierarchy consistency
extension Button {
    
    override func layoutSubviews() {
        super.layoutSubviews()
        
        // Make sure shadowView and backgroundView are always in the back
        [shadowView, backgroundView]
            .enumerated()
            .forEach { index, view in
                if view.superview?.subviews.firstIndex(of: view) != index {
                    view.removeFromSuperview()
                    view.frame = bounds
                    insertSubview(view, at: index)
                } else {
                    view.frame = bounds
                }
            }
    }
}

// MARK: - StateValueContainer getting and setting
extension Button {
    
    func stateValueContainer<Value>(for keyPath: KeyPath<Button, Value>) -> StateValueContainer<Value> {
        if let key = stateValueContainers.keys.first(where: { $0 == keyPath }),
           let container = stateValueContainers[key]?.base as? StateValueContainer<Value> {
            return container
        }
        let container = StateValueContainer<Value>()
        stateValueContainers[keyPath as PartialKeyPath] = AnyStateValueContainer(container)
        return container
    }

    @discardableResult
    func setValue<Value: Equatable>(_ value: Optional<Value>, for keyPath: KeyPath<Button, Value>, state: State) -> StateValueContainer<Value> {
        let container = stateValueContainer(for: keyPath)
        container.set(value, for: state)
        container.update(object: self, for: self.state, using: keyPath)
        return container
    }
    
    func value<Value: Equatable>(for keyPath: KeyPath<Button, Value>, state: State) -> Value? {
        stateValueContainer(for: keyPath).value(for: state)
    }
    
    func value<Value: Equatable>(for keyPath: KeyPath<Button, Optional<Value>>, state: UIButton.State) -> Value? {
        stateValueContainer(for: keyPath).value(for: state) as? Value
    }
    
    func processedValue<Value: Equatable>(for keyPath: KeyPath<Button, Value>, state: State, defaultValue: Value) -> (usedState: State, value: Value)? {
        stateValueContainer(for: keyPath).processedValue(for: state, defaultValue: defaultValue)
    }
    
    func processedValue<Value: Equatable>(for keyPath: KeyPath<Button, Optional<Value>>, state: State, defaultValue: Value? = nil) -> (usedState: State, value: Value?)? {
        stateValueContainer(for: keyPath).processedValue(for: state, defaultValue: defaultValue)
    }
}

// MARK: - Background color
extension Button {
    
    func setBackgroundColor(_ color: UIColor?, for state: State) {
        setValue(color, for: \.backgroundView.backgroundColor, state: state)
    }
    
    func backgroundColor(for state: State) -> UIColor? {
        value(for: \.backgroundView.backgroundColor, state: state)
    }
}

// MARK: - Image color
extension Button {
    
    override func setImage(_ image: UIImage?, for state: UIControl.State) {
        if image?.renderingMode != .alwaysTemplate, imageColor(for: state) != nil {
            super.setImage(image?.withRenderingMode(.alwaysTemplate), for: state)
            return
        }
        super.setImage(image, for: state)
    }
    
    func setImageColor(_ color: UIColor?, for state: State) {
        if let image = image(for: state) {
            setImage(image.withRenderingMode(.alwaysTemplate), for: state)
        }
        if color != nil {
            adjustsImageWhenHighlighted = false
            adjustsImageWhenDisabled = false
        }
        setValue(color, for: \.imageView?.tintColor, state: state)
            .onUpdate { [unowned self] state, container in
                self.adjustsImageWhenHighlighted = !container.isEmpty
                self.adjustsImageWhenDisabled = !container.isEmpty
                
                if let result = container.processedValue(for: state, defaultValue: self.originalTintColor) {
                    imageView?.tintColor = result.value
                }
            }
    }
    
    func imageColor(for state: State) -> UIColor? {
        value(for: \.imageView?.tintColor, state: state)
    }
}

// MARK: - Border color
extension Button {
    
    func setBorderColor(_ color: UIColor?, for state: State) {
        setValue(color?.cgColor, for: \.backgroundView.layer.borderColor, state: state)
    }
    
    func borderColor(for state: State) -> UIColor? {
        value(for: \.backgroundView.layer.borderColor, state: state).map { UIColor(cgColor: $0) }
    }
}

StateValueContainer.swift

import UIKit

extension UIControl.State: Hashable { }

protocol StateValueContainable {
    var animatesUpdate: Bool { get }
    func update<Object: AnyObject>(object: Object, for state: UIButton.State, using keyPath: PartialKeyPath<Object>)
}

struct AnyStateValueContainer: StateValueContainable {
    
    let base: Any
    
    init<Container>(_ wrapped: Container) where Container: StateValueContainable {
        base = wrapped
    }
    
    var animatesUpdate: Bool {
        (base as? StateValueContainable)?.animatesUpdate == true
    }
    
    func update<Object: AnyObject>(object: Object, for state: UIButton.State, using keyPath: PartialKeyPath<Object>) {
        guard let container = base as? StateValueContainable else {
            return
        }
        container.update(object: object, for: state, using: keyPath)
    }
}

/// Handles storing and retrieving values linked to button states
/// This allows for adding specific values to a UIButton linked
/// to any (combo of) button state, like fonts and colors
final class StateValueContainer<Value: Equatable>: StateValueContainable {
    
    private(set) var animatesUpdate: Bool = false
    
    private var storage = [UIButton.State: Value]()
    private var updater: ((UIButton.State, StateValueContainer<Value>) -> Void)?
    
    subscript(state: UIButton.State) -> Value? {
        get {
            return value(for: state)
        }
        set {
            set(newValue, for: state)
        }
    }
    
    /// Sets the value for the given state, removes it if `nil`
    func set(_ value: Value?, for state: UIButton.State) {
        guard let value = value else {
            storage.removeValue(forKey: state)
            return
        }
        storage[state] = value
    }
    
    /// Returns the value if available
    func value(for state: UIButton.State) -> Value? {
        storage[state]
    }
    
    /// Whether any values have been stored
    var isEmpty: Bool {
        return storage.isEmpty
    }
    
    /// To closely approximate `UIButton`'s default behavior of using the normal state as a fallback state
    /// and only use a default (or nil) value when there is no value set for the normal state, use this method
    /// when you need to update a specific value in your view
    ///
    /// - This method will return nil when there's no values set, as to not override any view properties
    /// set outside of the set/get methods.
    /// - Will either use the value for the state given or the value set for the `.normal` state
    /// - Only returns `defaultValue` when the requested state is `.normal`, otherwise the view's
    /// properties should not be changed
    ///
    /// - Parameters:
    ///   - state: Control state to request the value for
    ///   - defaultValue: Value to use when no value available for the given state or `.normal` state
    /// - Returns: Tuple with the used state and the value to set for the value
    func processedValue(for state: UIButton.State, defaultValue: Value? = nil) -> (usedState: UIButton.State, value: Value?)? {
        guard !isEmpty else {
            return nil
        }
        
        var result = (usedState: state, value: nil as Value?)
        
        if let value = value(for: state) {
            result.value = value
        } else if state == .normal || value(for: .normal) != nil {
            result.usedState = .normal
            result.value = value(for: .normal) ?? defaultValue
        }
        return result
    }
    
    /// To closely approximate `UIButton`'s default behavior of using the normal state as a fallback state
    /// and only use a default (or nil) value when there is no value set for the normal state, use this method
    /// when you need to update a specific value in your view
    ///
    /// - This method will return nil when there's no values set, as to not override any view properties
    /// set outside of the set/get methods.
    /// - Will either use the value for the state given or the value set for the `.normal` state
    /// - Only returns `defaultValue` when the requested state is `.normal`, otherwise the view's
    /// properties should not be changed
    ///
    /// - Note: `defaultValue` is not optional in this method to make sure to always get a value in the
    /// `.normal` state
    ///
    /// - Parameters:
    ///   - state: Control state to request the value for
    ///   - defaultValue: Value to use when no value available for the given state or `.normal` state
    /// - Returns: Tuple with the used state and the value to set for the value
    func processedValue(for state: UIButton.State, defaultValue: Value) -> (usedState: UIButton.State, value: Value)? {
        guard let result = processedValue(for: state, defaultValue: Optional(defaultValue)) as? (UIButton.State, Value) else {
            return nil
        }
        return result
    }
    
    func update<Object: AnyObject>(object: Object, for state: UIButton.State, using keyPath: PartialKeyPath<Object>) {
        if let updater = updater {
            updater(state, self)
            return
        }
        
        if let keyPath = keyPath as? ReferenceWritableKeyPath<Object, Value> {
            if let value = processedValue(for: state)?.value {
                object[keyPath: keyPath] = value
            }
        } else if let keyPath = keyPath as? ReferenceWritableKeyPath<Object, Optional<Value>> {
            if let result = processedValue(for: state) {
                object[keyPath: keyPath] = result.value
            }
        }
        
    }
}

extension StateValueContainer {
    @discardableResult
    func onUpdate(_ handler: @escaping (UIButton.State, StateValueContainer<Value>) -> Void) -> Self {
        updater = handler
        return self
    }
    
    @discardableResult
    func animatesUpdate(_ animates: Bool = true) -> Self {
        animatesUpdate = true
        return self
    }
}