Мавпа виправляє клас в іншому модулі на Python

Я працюю з модулем, написаним кимось іншим. Я хотів би, щоб мавпа виправила __init__метод класу, визначеного в модулі. Приклади, які я знайшов, що показують, як це зробити, припускали, що я сам закликаю клас (наприклад, клас Monkey-patch Python ). Однак це не так. У моєму випадку клас італізується в межах функції в іншому модулі. Дивіться (значно спрощений) приклад нижче:

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

Чи є спосіб модифікувати __init__метод, SomeClassщоб при dosomethingвиклику з mymodule.py він, наприклад, друкував 43 замість 42? В ідеалі я міг би обернути існуючий метод.

Я не можу змінити файли Third-PartyModule * .py, оскільки інші сценарії залежать від існуючих функціональних можливостей. Я волів би не створювати власну копію модуля, оскільки зміна, яку мені потрібно внести, дуже проста.

Редагувати 24.10.2013

Я пропустив невелику, але важливу деталь у наведеному вище прикладі. SomeClassімпортується thirdpartymodule_bяк це: from thirdpartymodule_a import SomeClass.

Щоб зробити патч, запропонований FJ, мені потрібно замінити копію thirdpartymodule_b, а не thirdpartymodule_a. напр thirdpartymodule_b.SomeClass.__init__ = new_init.

Має працювати наступне:

import thirdpartymodule_a
import thirdpartymodule_b

def new_init(self):
    self.a = 43

thirdpartymodule_a.SomeClass.__init__ = new_init

thirdpartymodule_b.dosomething()

Якщо ви хочете, щоб новий init викликав старий init, замініть new_init()визначення таким:

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

Використовувати 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

або

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

Брудно, але це працює:

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

Ще один можливий підхід, дуже схожий на підхід Ендрю Кларка , полягає у використанні бібліотеки обгортки . Крім іншого корисного, ця бібліотека надає wrap_function_wrapperі patch_function_wrapperпомічників. Їх можна використовувати так:

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

Або іноді вам може знадобитися використовувати wrap_function_wrapperне декоратор, а в іншому випадку працює однаково:

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

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

Одна лише трохи менш удача версія використовує глобальні змінні як параметри:

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

коли сторожовий неправда, він діє так само, як і раніше. Коли це правда, тоді ви отримуєте свою нову поведінку. У коді ви зробите:

import thirdpartymodule_b

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

Звичайно, досить тривіально зробити це належним чином, не впливаючи на існуючий код. Але вам доведеться трохи змінити інший модуль:

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

і перейдіть до 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

Існуючий код буде продовжувати працювати - вони будуть викликати його без аргументів, що збереже значення за замовчуванням false, яке збереже стару поведінку. Але ваш код тепер має спосіб повідомити цілому стеку про те, що доступна нова поведінка.

Ось приклад, який я придумав для monkeypatch Popenза допомогоюpytest .

імпортувати модуль:

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

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

MockPopenЗавод зібрати макет Попенс:

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

І приклад тесту:

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

Це той самий приклад, але pytest.fixtureйого використання замінює Popenімпорт вбудованого класу в межах 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']