Finding imports of a function globally

Question

I am working on a test fixture for the slash framework which needs to modify the behavior of time.sleep. For reasons I cannot use pytest, so I am trying to roll my own basic monkeypatching support.

I am able to replace time.sleep easily enough for things that just import time, but some things do from time import sleep before my fixture is instantiated. So far I'm using gc.get_referrers to track down any references to sleep before replacing them:

self._real_sleep = time.sleep
for ref in gc.get_referrers(time.sleep):
    if (isinstance(ref, dict)
            and "__name__" in ref
            and "sleep" in ref
            and ref["sleep"] is self._real_sleep):
        self._monkey_patch(ref)

In practice this works, but it feels very ugly. I do have access to reasonably current python (currently on 3.11), just limited ability to add 3rd party dependencies. Is there a better/safer way to find references to a thing or to patch out a thing globally, using only standard library methods?

Answer 1

YOu can use Python's sys.modules to replace time.sleep globally. This includes both direct imports of time.sleep as well as cases where from time import sleep is used.

Here's the code.

import time
import sys
import gc

class TimeSleepPatch:
    def __init__(self, new_sleep_func):
        self._new_sleep_func = new_sleep_func
        self._real_sleep = time.sleep
        self._patched_modules = []

    def _patch_module(self, module, attr_name):
        # Replace the reference to the original time.sleep with the patched one
        if hasattr(module, attr_name):
            if getattr(module, attr_name) is self._real_sleep:
                setattr(module, attr_name, self._new_sleep_func)
                self._patched_modules.append((module, attr_name))

    def patch(self):
        # Patch the time module itself
        time.sleep = self._new_sleep_func

        # Go through all modules and replace any reference to time.sleep
        for module in sys.modules.values():
            if module:
                self._patch_module(module, "sleep")

        # Optional: Patch referrers in case any obscure reference still lingers
        for ref in gc.get_referrers(self._real_sleep):
            if isinstance(ref, dict):
                for key, value in ref.items():
                    if value is self._real_sleep:
                        ref[key] = self._new_sleep_func

    def restore(self):
        # Restore the original time.sleep
        time.sleep = self._real_sleep

        # Restore patched modules to their original state
        for module, attr_name in self._patched_modules:
            setattr(module, attr_name, self._real_sleep)
        self._patched_modules.clear()

# Example of a custom sleep function
def fake_sleep(duration):
    print(f"Fake sleep for {duration} seconds")

# Instantiate the patcher and patch the time.sleep
patcher = TimeSleepPatch(fake_sleep)
patcher.patch()

# Test it (this will print "Fake sleep for 2 seconds")
time.sleep(2)

# Restore the original time.sleep
patcher.restore()

# Test the restored behavior (this will actually sleep)
time.sleep(2)

I hope this will help you a little.