iOS --- 为UIButton添加setBackgroundColor:forState:方法(包含OC和Swift两个版本)

有这样的一类简单需求: UIButton的背景色要与其state相关, 如未点击时显示蓝色, 点击时显示绿色.
但是, UIButton自身并未提供setBackgroundColor:forState:方法, 因此我们不得不单独在touchDown等方法中去更新其backgroundColor属性.
这里介绍如何为UIButton提供该扩展方法, Objective-C和Swift的版本都有.
其中用到了runtime的关联对象, 不熟悉的同学可以先参考iOS — 理解Runtime机制及其使用场景.

Objective-C

Objective-C中通过Category提供该扩展方法.
头文件:

#import <UIKit/UIKit.h>

@interface UIButton (CS_BackgroundColor)

- (void)setBackgroundColor:(UIColor *)backgroundColor forState:(UIControlState)state;

@end

实现文件:

#import "UIButton+CS_BackgroundColor.h"
#import <objc/runtime.h>

@interface UIButton (CS_BackgroundColor)

@property (nonatomic, strong) NSMutableDictionary *cs_dictBackgroundColor;

@end


@implementation UIButton (CS_BackgroundColor)

static const NSString *key_cs_backgroundColor             = @"key_cs_backgroundColor";

static NSString *cs_stringForUIControlStateNormal         = @"cs_stringForUIControlStateNormal";
static NSString *cs_stringForUIControlStateHighlighted    = @"cs_stringForUIControlStateHighlighted";
static NSString *cs_stringForUIControlStateDisabled       = @"cs_stringForUIControlStateDisabled";
static NSString *cs_stringForUIControlStateSelected       = @"cs_stringForUIControlStateSelected";

#pragma mark - cs_dictBackgroundColor

- (NSMutableDictionary *)cs_dictBackgroundColor {
    return objc_getAssociatedObject(self, &key_cs_backgroundColor);
}

- (void)setCs_dictBackgroundColor:(NSMutableDictionary *)cs_dictBackgroundColor {
    objc_setAssociatedObject(self, &key_cs_backgroundColor, cs_dictBackgroundColor, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
}

- (void)setBackgroundColor:(UIColor *)backgroundColor forState:(UIControlState)state {
    if (!self.cs_dictBackgroundColor) {
        self.cs_dictBackgroundColor = [[NSMutableDictionary alloc] init];
    }

    [self.cs_dictBackgroundColor setObject:backgroundColor forKey:[self cs_stringForUIControlState:state]];
}

- (NSString *)cs_stringForUIControlState:(UIControlState)state {
    NSString *cs_string;
    switch (state) {
        case UIControlStateNormal:
            cs_string = cs_stringForUIControlStateNormal;
            break;
        case UIControlStateHighlighted:
            cs_string = cs_stringForUIControlStateHighlighted;
            break;
        case UIControlStateDisabled:
            cs_string = cs_stringForUIControlStateDisabled;
            break;
        case UIControlStateSelected:
            cs_string = cs_stringForUIControlStateSelected;
            break;
        default:
            cs_string = cs_stringForUIControlStateNormal;
            break;
    }
    return cs_string;
}

#pragma mark - highlighted

- (void)setHighlighted:(BOOL)highlighted {
    if (highlighted) {
        self.backgroundColor = (UIColor *)[self.cs_dictBackgroundColor objectForKey:cs_stringForUIControlStateHighlighted];
    } else {
        self.backgroundColor = (UIColor *)[self.cs_dictBackgroundColor objectForKey:cs_stringForUIControlStateNormal];
    }
}

@end

使用方法:

btn1.backgroundColor = [UIColor greenColor]; // 注意, 这里的默认背景色必须设置, 仅仅通过下一行暂时不能设置初始的背景色
[btn1 setBackgroundColor:[UIColor greenColor] forState:UIControlStateNormal];
[btn1 setBackgroundColor:[UIColor blueColor] forState:UIControlStateHighlighted];

Swift

Swift的语法终归是要不断练习. 通过extension来提供该方法.

public extension UIButton {

    private struct cs_backgroundColor {
        static var keyBackgroundColors              = "cs_keyBackgroundColors"

        static var keyBackgroundColor_Normal        = "cs_keyBackgroundColor_Normal"
        static var keyBackgroundColor_Highlighted   = "cs_keyBackgroundColor_Highlighted"
    }

    var cs_dictBackgroundColors: Dictionary<String, UIColor>! {
        get {
            if let dictBackgroundColors = objc_getAssociatedObject(self, &cs_backgroundColor.keyBackgroundColors) {
                return dictBackgroundColors as! Dictionary<String, UIColor>
            }

            return nil
        }

        set {
            objc_setAssociatedObject(self, &cs_backgroundColor.keyBackgroundColors, newValue as Dictionary<String, UIColor>, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
        }
    }

    func cs_setBackgroundColor(color: UIColor, forState: UIControlState) {
        if self.cs_dictBackgroundColors == nil {
            self.cs_dictBackgroundColors = Dictionary<String, UIColor>()
        }

        if let key = self.cs_stringForUIControlState(forState) {
            self.cs_dictBackgroundColors[key] = color
        }
    }

    private func cs_stringForUIControlState(state: UIControlState) -> String! {
        var cs_string = ""

        switch state {
        case UIControlState.Normal:
            cs_string = cs_backgroundColor.keyBackgroundColor_Normal
        case UIControlState.Highlighted:
            cs_string = cs_backgroundColor.keyBackgroundColor_Highlighted
        default:
            cs_string = cs_backgroundColor.keyBackgroundColor_Normal
        }

        return cs_string
    }

    override var highlighted: Bool {
        get {
            return super.highlighted
        }

        set {
            if newValue {
                if let key = self.cs_stringForUIControlState(.Highlighted) {
                    self.backgroundColor = self.cs_dictBackgroundColors[key]!
                }
            } else {
                if let key = self.cs_stringForUIControlState(.Normal) {
                    self.backgroundColor = self.cs_dictBackgroundColors[key]!
                }
            }
        }
    }

}

使用方法:

btn.backgroundColor = UIColor.greenColor() // 这里的设置同样是必需的.
btn.cs_setBackgroundColor(UIColor.greenColor(), forState: .Normal)
btn.cs_setBackgroundColor(UIColor.blueColor(), forState: .Highlighted)

Demo

Objective-C版本的Demo请参考DemoRuntime.
Swift版本的Demo请参考CSSwiftExtension.