Python monkey patching: instance creation in method of library/object

Question

what is the easiest way to solve the following problem in extending/altering the functionality of a third party library?

The library offers a class LibraryClass with a function func_to_be_changed. This function has a local variable internal_variable which is the instance of another class SimpleCalculation of that library. I created a new class BetterCalculation in my own module and now want LibraryClass.func_to_be_changed to use an instance of this new class.

# third party library
from third_party_library.utils import SimpleCalculation

class LibraryClass:
    def func_to_be_changed(self, x):
        # many complicated things go on
        internal_variable = SimpleCalculation(x)
        # many more complicated things go on

The easiest solution would be, to just copy the code from the third party library, subclass the LibraryClass and overwrite the function func_to_be_changed:

# my module
from third_party_library import LibraryClass

class BetterLibraryClass(LibraryClass):
    def func_to_be_changed(self, x):
        """This is an exact copy of LibraryClass.func_to_be_changed."""
        # many complicated things go on
        internal_variable = BetterCalculation(x)  # Attention: this line has been changed!!!
        # many more complicated things go on

However, this involves copying of many lines of code. When a new version of the third party class improves on code that was copied without modification, this modifications need to be incorporated manually by another copying step.

I tried to use unittest.mock.patch as I know that the following two snippets work:

# some script
from unittest.mock import patch
import third_party_library

from my_module import BetterCalculation

with patch('third_party_library.utils.SimpleCalculation', BetterCalculation):
    local_ = third_party_library.utils.SimpleCalculation(x)  # indeed uses `BetterCalculation`


def foo(x):
    return third_party_library.utils.SimpleCalculation(x)


with patch('third_party_library.utils.SimpleCalculation', BetterCalculation):
    local_ = foo(x)  #  indeed uses `BetterCalculation`

However, the following does not work:

# some script
from unittest.mock import patch
from third_party_library.utils import SimpleCalculation

from my_module import BetterCalculation

def foo(x):
    return SimpleCalculation(x)

with patch('third_party_library.utils.SimpleCalculation', BetterCalculation):
    local_ = foo(x)  # does not use `BetterCalculation`

# this works again
with patch('__main__.SimpleCalculation', BetterCalculation):
    local_ = foo(x)  #  indeed uses `BetterCalculation`

Therefore, the following won`t work either:

# my module
from unittest.mock import patch
from third_party_library import LibraryClass

from my_module import BetterCalculation

class BetterLibraryClass(LibraryClass):
    def func_to_be_changed(self, x):
        with patch(
            'third_party_library.utils.SimpleCalculation',
            BetterCalculation
        ):
            super().func_to_be_changed(x)

Is there an elegant pythonic way to do this? I guess this boils down to the question: What is the equivaltent of __main__ in the last code snippet that needs to replace third_party_library.utils?

Answer 1

Some context

The first string argument in the patch function can have two different meanings depending on the situation. In the first situation the described object has not been imported and is unavailable to the program which would, therefore, result in a NameError without the mocking. However, in the question, an object needs to be overwritten. Therefore, it is available to the program and not using patch would not result in an error.

Disclaimer

I might have used the complete wrong language in here, as for sure there are precise python terms for all the described notions.

Overwriting an object

As shown in the question, the locally imported SimpleCalculation can be overwritten with __main__.SimpleCalculation. Therefore, it is important to remember that you need to tell patch the path to the local object and not how that same object would be imported in the current script. Let's assume the following module:

# thirdpartymodule/__init__.py
from .utils import foo


def local_foo():
    print("Hello local!")


class Bar:
    def __init__(self):
        foo()
        local_foo()

and

# thirdpartymodule/utils.py
def foo():
    print("third party module")

We want to override the functions foo and local_foo. But we don't want to override any functions, we want to override the functions foo and local_foo in the context of the file thirdpartymodule/__init__.py. It is unimportant that the function foo enters the context of the file via an import statement, while local_foo is defined locally. So we want to override the functions in the context of thirdpartymodule.foo and thirdpartymodule.local_foo. The context thirdpartymodule.utils.foo is not important here and won't help us. The following snippet illustrates that:

from unittest.mock import patch
from thirdpartymodule import Bar


bar = Bar()
# third party module
# Hello local!

def myfoo():
    print("patched function")
    
    
with patch("thirdpartymodule.foo", myfoo):
    bar = Bar()
    # patched function
    # Hello local!
    
# will not work!
with patch("thirdpartymodule.utils.foo", myfoo):
    bar = Bar()
    # third party module
    # Hello local!
    
with patch("thirdpartymodule.local_foo", myfoo):
    bar = Bar()
    # third party module
    # patched function

In the hypothetical module of the question we first need to assume that the class LibraryClass is defined in the file third_party_library/library_class.py. Then, we want to override third_party_library.library_class.SimpleCalculation and the correct patch would be:

# my module
from unittest.mock import patch
from third_party_library import LibraryClass

from my_module import BetterCalculation

class BetterLibraryClass(LibraryClass):
    def func_to_be_changed(self, x):
        with patch(
            'third_party_library.library_class.SimpleCalculation',
            BetterCalculation
        ):
            super().func_to_be_changed(x)