Monkey 在 Python 的另一个模块中修补类

小开

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()

小开

最佳答案

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

小开

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

小开

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

sentinel = False


class SomeClass(object):
def __init__(self):
global sentinel
if sentinel:
<do my custom code>
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:
<do my custom code>
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.

小开

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']

小开

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)