Search code examples
pythonpython-3.xmockingmonkeypatching

Patch module entry and exit - Python


I have a module called my_module whose performance I want to measure.

my_module has multiple functions and what I want to do is, whenever any function from the module is called

  • Note down the entry time
  • Note down the exit time

So let's say I ran this - my_module.func('blah')

I wish to print

  • Function call started at 100001201
  • Function call ended at 100001312

for any function in the module.

I cannot make any direct changes to the module file as it has to be used in other places as well. I just need this to print for the performance test script.

Please help.


Solution

  • I did it on Numpy module, so you should be able to use the same method for for any module:

    import numpy as np                           # import the module you want to be logged
    from inspect import getmembers, isfunction   # These will give us the functions in a module
    import time                                  # for logging the time
    
    class Object(object):                        # make a dummy class on which we can do "setattr()"
        pass
    
    numpy_with_time_logger = Object()            # We are going to make numpy_with_time_logger work exactly like Numpy except that it will also print start and end times
    
    
    def return_call_time_logger(f):              # This function takes in a function and returns a new function!
        def call_time_logger(*args, **kargs):    # This function is going to substitute f. It runs f(), but it also prints the times.
            print(f'Function call started at {time.time()}')
            ret_val = f(*args, **kargs)
            print(f'Function call ended at {time.time()}')
            return ret_val
        return call_time_logger
    
    for f_name, f in getmembers(numpy, isfunction): # For every function inside numpy module
        # Set numpy_with_time_logger.f_name to be return_call_time_logger(f)
        setattr(numpy_with_time_logger, f_name, return_call_time_logger(f))
    

    Note: To make it work for your own module, just change the "numpy" in getmembers(numpy, isfunction) to the name of your imported module.

    Now, our numpy_with_time_logger object is set to work exactly like Numpy, but it will also print the start and end times. So try using any of numpy's functions through numpy_with_time_logger:

    print(numpy_with_time_logger.asarray([1,2]))
    
    

    The output is:

    start time 1615096471.112301
    end time 1615096471.1134446
    array([1, 2])
    

    An issue with this method:

    If you call numpy_with_time_logger.func1() and numpy_with_time_logger.func1() calls numpy.func2(), then you will only get the time logs for func1, and you won't get the time logs for func2.

    To access the global variables defined in Numpy module, just keep using numpy, e.g. numpy.int32 .