web-dev-qa-db-fra.com

Monkey patche une classe dans un autre module en Python

Je travaille avec un module écrit par quelqu'un d'autre. Je voudrais patcher le singe __init__ méthode d'une classe définie dans le module. Les exemples que j'ai trouvés montrant comment procéder ont tous supposé que j'appellerais la classe moi-même (par exemple Monkey-patch Python class ). Cependant, ceci est pas le cas. Dans mon cas, la classe est initalisée dans une fonction d'un autre module. Voir l'exemple (très simplifié) ci-dessous:

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

Existe-t-il un moyen de modifier le __init__ méthode de SomeClass de sorte que lorsque dosomething est appelée depuis mymodule.py, par exemple, elle affiche 43 au lieu de 42? Idéalement, je serais capable d'envelopper la méthode existante.

Je ne peux pas modifier les fichiers Thirdpartymodule * .py, car les autres scripts dépendent des fonctionnalités existantes. Je préfère ne pas avoir à créer ma propre copie du module, car la modification que je dois apporter est très simple.

Modifier le 2013-10-24

J'ai négligé un détail petit mais important dans l'exemple ci-dessus. SomeClass est importé par thirdpartymodule_b comme ça: from thirdpartymodule_a import SomeClass.

Pour faire le patch suggéré par F.J je dois remplacer la copie dans thirdpartymodule_b, plutôt que thirdpartymodule_a. par exemple. thirdpartymodule_b.SomeClass.__init__ = new_init.

pythonclassunit-testingmonkeypatching
58
Snorfalorpagus

Les éléments suivants devraient fonctionner:

import thirdpartymodule_a
import thirdpartymodule_b

def new_init(self):
    self.a = 43

thirdpartymodule_a.SomeClass.__init__ = new_init

thirdpartymodule_b.dosomething()

Si vous voulez que le nouveau init appelle l'ancien init, remplacez la définition new_init() par ce qui suit:

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

Utilisez la bibliothèque mock .

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

ou

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

Une autre approche possible, très similaire à celle d'Andrew Clark , consiste à utiliser la bibliothèque wrapt . Entre autres choses utiles, cette bibliothèque fournit wrap_function_wrapper et patch_function_wrapper aides. Ils peuvent être utilisés comme ceci:

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

Ou parfois, vous pouvez utiliser wrap_function_wrapper qui n'est pas décorateur mais othrewise fonctionne de la même manière:

def new_init(wrapped, instance, args, kwargs):
    pass  # ...

wrapt.wrap_function_wrapper(thirdpartymodule_a.SomeClass, '__init__', new_init)
2
MarSoft

Une seule version légèrement moins hacky utilise des variables globales comme paramètres:

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

quand la sentinelle est fausse, elle agit exactement comme avant. Quand c'est vrai, alors vous obtenez votre nouveau comportement. Dans votre code, vous feriez:

import thirdpartymodule_b

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

Bien sûr, il est assez trivial d'en faire un correctif correct sans impact sur le code existant. Mais vous devez changer légèrement l'autre module:

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

et passez à 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

Le code existant continuera de fonctionner - ils l'appelleront sans argument, ce qui conservera la valeur fausse par défaut, ce qui conservera l'ancien comportement. Mais votre code a maintenant un moyen de dire à la pile entière que de nouveaux comportements sont disponibles.

1
Corley Brigman

Sale, mais ça marche:

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
1
lucasg

Voici un exemple que j'ai trouvé pour monkeypatch Popen en utilisant pytest.

importer le module:

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

Un objet MockBytes:

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

Une usine MockPopen pour collecter les faux 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

Et un exemple de 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']

C'est le même exemple, mais en utilisant pytest.fixture il remplace l'importation de classe Popen intégrée dans 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']
1
jmunsch