添加使用 block 交换方法实现的接口

rentzsch · Nov 28, 2016 · 46c69cc · 46c69cc
1 parent d895bff
commit 46c69cc
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 0 deletions.
diff --git a/JRSwizzle.h b/JRSwizzle.h
@@ -10,4 +10,39 @@
 + (BOOL)jr_swizzleMethod:(SEL)origSel_ withMethod:(SEL)altSel_ error:(NSError**)error_;
 + (BOOL)jr_swizzleClassMethod:(SEL)origSel_ withClassMethod:(SEL)altSel_ error:(NSError**)error_;
 
+
+/**
+ ```
+ __block NSInvocation *invocation = nil;
+ invocation = [self jr_swizzleMethod:@selector(initWithCoder:) withBlock:^(id obj, NSCoder *coder) {
+ NSLog(@"before %@, coder %@", obj, coder);
+
+ [invocation setArgument:&coder atIndex:2];
+ [invocation invokeWithTarget:obj];
+
+ id ret = nil;
+ [invocation getReturnValue:&ret];
+
+ NSLog(@"after %@, coder %@", obj, coder);
+
+ return ret;
+ } error:nil];
+ ```
+ */
++ (NSInvocation*)jr_swizzleMethod:(SEL)origSel withBlock:(id)block error:(NSError**)error;
+
+/**
+ ```
+ __block NSInvocation *classInvocation = nil;
+ classInvocation = [self jr_swizzleClassMethod:@selector(test) withBlock:^() {
+ NSLog(@"before");
+
+ [classInvocation invoke];
+
+ NSLog(@"after");
+ } error:nil];
+ ```
+ */
++ (NSInvocation*)jr_swizzleClassMethod:(SEL)origSel withBlock:(id)block error:(NSError**)error;
+
 @end
diff --git a/JRSwizzle.m b/JRSwizzle.m
@@ -139,4 +139,28 @@ + (BOOL)jr_swizzleClassMethod:(SEL)origSel_ withClassMethod:(SEL)altSel_ error:(
 	return [GetClass((id)self) jr_swizzleMethod:origSel_ withMethod:altSel_ error:error_];
 }
 
++ (NSInvocation*)jr_swizzleMethod:(SEL)origSel withBlock:(id)block error:(NSError**)error {
+    IMP blockIMP = imp_implementationWithBlock(block);
+    NSString *blockSelectorString = [NSString stringWithFormat:@"_jr_block_%@_%p", NSStringFromSelector(origSel), block];
+    SEL blockSel = sel_registerName([blockSelectorString cStringUsingEncoding:NSUTF8StringEncoding]);
+    Method origSelMethod = class_getInstanceMethod(self, origSel);
+    const char* origSelMethodArgs = method_getTypeEncoding(origSelMethod);
+    class_addMethod(self, blockSel, blockIMP, origSelMethodArgs);
+
+    NSMethodSignature *origSig = [NSMethodSignature signatureWithObjCTypes:origSelMethodArgs];
+    NSInvocation *origInvocation = [NSInvocation invocationWithMethodSignature:origSig];
+    origInvocation.selector = blockSel;
+
+    [self jr_swizzleMethod:origSel withMethod:blockSel error:nil];
+
+    return origInvocation;
+}
+
++ (NSInvocation*)jr_swizzleClassMethod:(SEL)origSel withBlock:(id)block error:(NSError**)error {
+    NSInvocation *invocation = [GetClass((id)self) jr_swizzleMethod:origSel withBlock:block error:error];
+    invocation.target = self;
+
+    return invocation;
+}
+
 @end