Objective-Cにて動的にメソッドを書き換える君を作ってみた（車輪の再発明）

OCMockではクラスメソッドは書き換えられるが、全インスタンスに影響するインスタンスメソッドを書き換えられなさそうなので書いてみた。 OCMockでクラスにかかるインスタンスメソッドを変更する方法を調べるのに半日ほど無駄にしてしまったので、何も形を残さないと本当に無駄になってしまいそうだったから作ってみた。

機能

• クラスメソッド・インスタンスメソッドを任意のblockで置き換え
• blockを書くの面倒だから、返り値が定数なら定数を返すだけ版も用意

注意

• エラー処理あんまり考えていないｗ
• 元に戻す処理は行っていない
• 同じメソッドに対して複数回呼び出すと、最後のやつが勝つ

これ

#import <objc/runtime.h>

@interface NSObject (TestMethodReplacer)
+ (void)replaceClassMethodWithSelector:(SEL)selector block:(id)block;
+ (void)replaceInstanceMethodWithSelector:(SEL)selector block:(id)block;

// convenience
+ (void)replaceClassMethodWithSelector:(SEL)selector returnObject:(id)returnObject;
+ (void)replaceInstanceMethodWithSelector:(SEL)selector returnObject:(id)returnObject;
@end

@implementation NSObject (TestMethodReplacer)
+ (void)replaceClassMethodWithSelector:(SEL)selector block:(id)block
{
    if (!class_respondsToSelector(self, selector)) return;
    [self _replaceMethod:selector block:block method:class_getClassMethod(self, selector)];
}

+ (void)replaceInstanceMethodWithSelector:(SEL)selector block:(id)block
{
    if (![self instancesRespondToSelector:selector]) return;
    [self _replaceMethod:selector block:block method:class_getInstanceMethod(self, selector)];
}

+ (void)_replaceMethod:(SEL)selector block:(id)block method:(Method)method
{
    struct objc_method_description *desc = method_getDescription(method);
    IMP imp = imp_implementationWithBlock(block);

    class_replaceMethod(self, selector, imp, desc->types);
}

+ (void)replaceClassMethodWithSelector:(SEL)selector returnObject:(id)returnObject
{
    [self replaceClassMethodWithSelector:selector block:^{
        return returnObject;
    }];
}

+ (void)replaceInstanceMethodWithSelector:(SEL)selector returnObject:(id)returnObject
{
    [self replaceInstanceMethodWithSelector:selector block:^{
        return returnObject;
    }];
}
@end

こんな感じで使う。

@interface Hoge : NSObject
+ (NSString *)hoge:(NSString *)str;
- (NSString *)fuga:(NSString *)str;
@end

@implementation Hoge
+ (NSString *)hoge:(NSString *)str
{
    return [NSString stringWithFormat:@"hoge, %@", str];
}

- (NSString *)fuga:(NSString *)str
{
    return [NSString stringWithFormat:@"fuga, %@", str];
}
@end

//////////

[Hoge hoge:@"Yo"];        //=> hoge, Yo
[[Hoge new] fuga:@"Yo"];  //=> fuga, Yo

// hoge: を書き換える
[Hoge replaceClassMethodWithSelector:@selector(hoge:) returnObject:@"dummy"];

[Hoge hoge:@"Yo"];        //=> dummy
[[Hoge new] fuga:@"Yo"];  //=> fuga, Yo

// fuga: を書き換える
[Hoge replaceInstanceMethodWithSelector:@selector(fuga:) returnObject:@"homu!"];

[Hoge hoge:@"Yo"];        //=> dummy
[[Hoge new] fuga:@"Yo"];  //=> homu!

// hoge: を書き換える
[Hoge replaceClassMethodWithSelector:@selector(hoge:) block: ^(id class, NSString *str){
    return [NSString stringWithFormat:@"dummy, %@", str];
}];

[Hoge hoge:@"Yo"];        //=> dummy, Yo
[[Hoge new] fuga:@"Yo"];  //=> homu!


// fuga: を書き換える
[Hoge replaceInstanceMethodWithSelector:@selector(fuga:) block: ^(id obj, NSString *str){
    return [NSString stringWithFormat:@"homu! %@", str];
}];

[Hoge hoge:@"Yo"];        //=> dummy, Yo
[[Hoge new] fuga:@"Yo"];  //=> homu! Yo

(面倒だったから例は実行確認していない)

感想とか

黒魔術楽しい。

今後は、restoreする仕組みを用意したい。 guardっぽいのは難しそう（というか、呼び出し側でautoreleasepoolで括らないと難しそう）。