Monkey patching a class in another module in Python

Question

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_b.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.

Sebaceous

I don't see why it would make a difference where you're calling the class from. – Slobber

Filenames should be thirdpartymodule_a.py, thirdpartymodule_b.py. – Promycelium

Answer 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

Endothermic

Maybe you should include a call to the old __init__. – Samira

Seems like inheriting from SomeClass and replacing the class would be much more elegant than messing with the __init__ functions themselves. – Pollywog

@JonathonReinhart You are probably right but I don't think the OP really wants to replace 42 by 43 in his own code. He specifically asked about monkey-patching though – Confound

Does this not create an infinite recursion if the new_init calls the old_init ? Because then new_init is assigned to old_init ? – Frager

@HakanBaba, no. Notice that we first save the original value of SomeClass.__init__ into a variable named old_init, then we overwrite SomeClass.__init__ with our implementation, and our implementation calls the saved old_init function. Infinite recursion would happen if you call SomeClass.__init__ from the new implementation. – Decolorize

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

Promycelium

this is the only way that actually works correctly. This basically monkey-patches while you do your call, does something, and then undoes the monkey-patch. That way, other modules calling it still get the original behaviour; only you get the modified behavior. (and thanks for pointing out mock!) – Mirabella

@CorleyBrigman This only applies to other modules within the same process. To me "other scripts" sounds like they would be separate Python processes that wouldn't be affected by a naive monkey-patch. – Endothermic

Answer 3

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)

Decolorize

Answer 4

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

Psi

what about if i want my new class definition extend the old one (via inheritance) ? – Linolinocut

why not inherit from SomeClass? – Kultur

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

Parrott

Answer 6

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.

Mirabella