Monkey patching a class in another module in Python

问题:

I'm working with a module written by someone else. I'd like to monkey patch the __init__ method of a class defined in the module. The examples I have found showing how to do this have all assumed I'd be calling the class myself (e.g. Monkey-patch Python class). However, this is not the case. In my case the class is initalised within a function in another module. See the (greatly simplified) example below:

thirdpartymodule_a.py

class SomeClass(object):     def __init__(self):         self.a = 42     def show(self):         print self.a 

thirdpartymodule_b.py

import thirdpartymodule_a def dosomething():     sc = thirdpartymodule_a.SomeClass()     sc.show() 

mymodule.py

import thirdpartymodule_b thirdpartymodule.dosomething() 

Is there any way to modify the __init__ method of SomeClass so that when dosomething is called from mymodule.py it, for example, prints 43 instead of 42? Ideally I'd be able to wrap the existing method.

I can't change the thirdpartymodule*.py files, as other scripts depend on the existing functionality. I'd rather not have to create my own copy of the module, as the change I need to make is very simple.

Edit 2013-10-24

I overlooked a small but important detail in the example above. SomeClass is imported by thirdpartymodule_b like this: from thirdpartymodule_a import SomeClass.

To do the patch suggested by F.J I need to replace the copy in thirdpartymodule_b, rather than thirdpartymodule_a. e.g. thirdpartymodule_b.SomeClass.__init__ = new_init.

回答1:

The following should work:

import thirdpartymodule_a import thirdpartymodule_b  def new_init(self):     self.a = 43  thirdpartymodule_a.SomeClass.__init__ = new_init  thirdpartymodule_b.dosomething() 

If you want the new init to call the old init replace the new_init() definition with the following:

old_init = thirdpartymodule_a.SomeClass.__init__ def new_init(self, *k, **kw):     old_init(self, *k, **kw)     self.a = 43 

回答2:

Use mock library.

import thirdpartymodule_a import thirdpartymodule_b import mock  def new_init(self):     self.a = 43  with mock.patch.object(thirdpartymodule_a.SomeClass, '__init__', new_init):     thirdpartymodule_b.dosomething() # -> print 43 thirdpartymodule_b.dosomething() # -> print 42 

or

import thirdpartymodule_b import mock  def new_init(self):     self.a = 43  with mock.patch('thirdpartymodule_a.SomeClass.__init__', new_init):     thirdpartymodule_b.dosomething() thirdpartymodule_b.dosomething() 

回答3:

Dirty, but it works :

class SomeClass2(object):     def __init__(self):         self.a = 43     def show(self):         print self.a  import thirdpartymodule_b  # Monkey patch the class thirdpartymodule_b.thirdpartymodule_a.SomeClass = SomeClass2  thirdpartymodule_b.dosomething() # output 43 

回答4:

One only slightly-less-hacky version uses global variables as parameters:

sentinel = False  class SomeClass(object):     def __init__(self):         global sentinel         if sentinel:                      else:             # Original code             self.a = 42     def show(self):         print self.a 

when sentinel is false, it acts exactly as before. When it's true, then you get your new behaviour. In your code, you would do:

import thirdpartymodule_b  thirdpartymodule_b.sentinel = True     thirdpartymodule.dosomething() thirdpartymodule_b.sentinel = False 

Of course, it is fairly trivial to make this a proper fix without impacting existing code. But you have to change the other module slightly:

import thirdpartymodule_a def dosomething(sentinel = False):     sc = thirdpartymodule_a.SomeClass(sentinel)     sc.show() 

and pass to init:

class SomeClass(object):     def __init__(self, sentinel=False):         if sentinel:                      else:             # Original code             self.a = 42     def show(self):         print self.a 

Existing code will continue to work - they will call it with no arguments, which will keep the default false value, which will keep the old behaviour. But your code now has a way to tell the whole stack on down that new behaviour is available.

回答5:

Here is an example I came up with to monkeypatch Popen using pytest.

import the module:

# must be at module level in order to affect the test function context from some_module import helpers 

A MockBytes object:

class MockBytes(object):      all_read = []     all_write = []     all_close = []      def read(self, *args, **kwargs):         # print('read', args, kwargs, dir(self))         self.all_read.append((self, args, kwargs))      def write(self, *args, **kwargs):         # print('wrote', args, kwargs)         self.all_write.append((self, args, kwargs))      def close(self, *args, **kwargs):         # print('closed', self, args, kwargs)         self.all_close.append((self, args, kwargs))      def get_all_mock_bytes(self):         return self.all_read, self.all_write, self.all_close 

A MockPopen factory to collect the mock popens:

def mock_popen_factory():     all_popens = []      class MockPopen(object):          def __init__(self, args, stdout=None, stdin=None, stderr=None):             all_popens.append(self)             self.args = args             self.byte_collection = MockBytes()             self.stdin = self.byte_collection             self.stdout = self.byte_collection             self.stderr = self.byte_collection             pass      return MockPopen, all_popens 

And an example test:

def test_copy_file_to_docker():     MockPopen, all_opens = mock_popen_factory()     helpers.Popen = MockPopen # replace builtin Popen with the MockPopen     result = copy_file_to_docker('asdf', 'asdf')     collected_popen = all_popens.pop()     mock_read, mock_write, mock_close = collected_popen.byte_collection.get_all_mock_bytes()     assert mock_read     assert result.args == ['docker', 'cp', 'asdf', 'some_container:asdf'] 

This is the same example, but using pytest.fixture it overrides the builtin Popen class import within helpers:

@pytest.fixture def all_popens(monkeypatch): # monkeypatch is magically injected      all_popens = []      class MockPopen(object):         def __init__(self, args, stdout=None, stdin=None, stderr=None):             all_popens.append(self)             self.args = args             self.byte_collection = MockBytes()             self.stdin = self.byte_collection             self.stdout = self.byte_collection             self.stderr = self.byte_collection             pass     monkeypatch.setattr(helpers, 'Popen', MockPopen)      return all_popens   def test_copy_file_to_docker(all_popens):         result = copy_file_to_docker('asdf', 'asdf')     collected_popen = all_popens.pop()     mock_read, mock_write, mock_close = collected_popen.byte_collection.get_all_mock_bytes()     assert mock_read     assert result.args == ['docker', 'cp', 'asdf', 'fastload_cont:asdf'] 

回答6:

One another possible approach, very similar to Andrew Clark's one, is to use wrapt library. Among other useful things, this library provides wrap_function_wrapper and patch_function_wrapper helpers. They can be used like this:

import wrapt import thirdpartymodule_a import thirdpartymodule_b  @wrapt.patch_function_wrapper(thirdpartymodule_a.SomeClass, '__init__') def new_init(wrapped, instance, args, kwargs):     # here, wrapped is the original __init__,     # instance is `self` instance (it is not true for classmethods though),     # args and kwargs are tuple and dict respectively.      # first call original init     wrapped(*args, **kwargs)  # note it is already bound to the instance     # and now do our changes     instance.a = 43  thirdpartymodule_b.do_something() 

Or sometimes you may want to use wrap_function_wrapper which is not a decorator but othrewise works the same way:

def new_init(wrapped, instance, args, kwargs):     pass  # ...  wrapt.wrap_function_wrapper(thirdpartymodule_a.SomeClass, '__init__', new_init) 

文章来源: Monkey patching a class in another module in Python