scoped_block_swizzler.mm 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. // Copyright 2013 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "ios/testing/scoped_block_swizzler.h"
  5. #include "base/check.h"
  6. #if !defined(__has_feature) || !__has_feature(objc_arc)
  7. #error "This file requires ARC support."
  8. #endif
  9. ScopedBlockSwizzler::ScopedBlockSwizzler(Class target, SEL selector, id block) {
  10. method_ = class_getInstanceMethod(target, selector);
  11. if (!method_) {
  12. // Try swizzling a class method instead.
  13. method_ = class_getClassMethod(target, selector);
  14. }
  15. DCHECK(method_);
  16. IMP block_imp = imp_implementationWithBlock(block);
  17. original_imp_ = method_setImplementation(method_, block_imp);
  18. }
  19. ScopedBlockSwizzler::ScopedBlockSwizzler(Class target,
  20. SEL selector,
  21. id block,
  22. BOOL class_method) {
  23. if (class_method) {
  24. method_ = class_getClassMethod(target, selector);
  25. } else {
  26. method_ = class_getInstanceMethod(target, selector);
  27. }
  28. DCHECK(method_);
  29. IMP block_imp = imp_implementationWithBlock(block);
  30. original_imp_ = method_setImplementation(method_, block_imp);
  31. }
  32. ScopedBlockSwizzler::~ScopedBlockSwizzler() {
  33. IMP block_imp = method_setImplementation(method_, original_imp_);
  34. DCHECK(imp_removeBlock(block_imp));
  35. }