Overriding functions imported by a subsequent import?

For a project, I need to substitute a function from a module. I would like to achieve this by the following minimal example

In mod.py the original function is defined,

def f():
    print("mod.py")

which should be overridden by f2 in mod2.py,

import sys

def f2():
    print("mod2.py")

sys.modules["mod"].f = f2

In main.py I first load the original function, which I immediately patch

from mod import f
import mod2

f()  # I'm expecting "mod2.py", but get "mod.py" instead

It seems, that the function definition of f in the __main__'s globals() is a not a reference to sys.modules["mod"].f, which I was hoping.

If in main.py I import the module in which f is defined, instead f itself, I can patch it

import mod
import mod2

print(mod.f())  # patched from mod2

But I'd like to find a solution which works both imports in main.py — the import of the whole module (import mod) and the function itself (from mod import f).

I also tried to somehow modify the locals() from the parent frame via inspect.currentframe().f_back.f_locals but this seems doesn't work as well and would feel hacky, if I'd find a solution like this, which would do what I want.

Is this in principle possible with the way Python handles imports?

Solution

Here's a good way to monkey-patch the function in mod.py so that it will affect any further imports of mod after an import mod2. This works because modules are cached in sys.modules — but you don't need to explicitly reference that if you do things as shown.

mod.py:

def f():
    print("mod.py")

mod2.py:

import mod

def f2():
    print("mod2.py")

# Monkey-patch `mod` module.
mod.f = f2

main.py:

import mod
import mod2  # Importing this applies the patch to the `mod` module.

mod.f()  # Prints -> mod2.py