Search code examples
pythonloggingdynamicmetaprogrammingwrapper

How to wrap python instance methods at runtime


I want to log debug information when calls are made in some class. The logged data is:

  • function name
  • stack trace of where the function was called from
  • time it took the function to execute
  • args and kwargs passed to the function

I want the wrapper to be fairly generic. At the same time, the wrapping should take place at run time. I've made the following wrapper class that logs the information and delegates the call to the original instance.

import datetime
import traceback
from functools import partial

from logging_module import db_log


class BaseWrapper(object):
    def __init__(self, item):
        self._item = item

    def __getattr__(self, attr):
        return getattr(self._item, attr)


class DBLogWrapper(BaseWrapper):

    @staticmethod
    def _time_method(method):
        name = "{0}.{1}.{2}".format(
            method.im_class.__module__,
            method.im_class.__name__,
            method.__name__
        )

        def timed_method(self, *args, **kwargs):
            begin = datetime.datetime.now()
            return_val = method.im_func(self, *args, **kwargs)
            end = datetime.datetime.now()

            trace = traceback.format_stack()

            db_log(
                name,
                begin,
                end,
                info={
                    'args': args,
                    'kwargs': kwargs,
                    'trace': trace
                }
            )
            return return_val

        return timed_method

    def __init__(self, item, methods):
        super(DBLogWrapper, self).__init__(item)
        for method in methods:
            class_method = getattr(item, method)
            wrapped_method = DBLogWrapper._time_method(class_method)
            wrapped_method = partial(wrapped_method, self._item)
            setattr(self, method, wrapped_method)

Example Usage:

class MyClass(object):

    def hello(self, greeting):
        print greeting

    def goodbye(self):
        print 'Good Bye'

a = MyClass()

if DEBUG:
    a = DBLogWrapper(a, ['hello'])

a.hello()
a.goodbye()

In this case the call to hello would be logged, but the call to goodbye would not.

However, this seems to be overkill for a task which seems like it should be simple. I'm looking for suggestions on how to improve the above code or an entirely different way to go about it.


Solution

  • You are doing too much work. You don't need partial at all. Just define timed_method without the self parameter and call method directly:

    import datetime
    import traceback
    from functools import partial
    
    def db_log(*args, **kwargs): print args, kwargs # Mock
    
    
    class BaseWrapper(object):
        def __init__(self, instance):
            self._instance = instance
    
        def __getattr__(self, attr):
            return getattr(self._instance, attr)
    
    
    class DBLogWrapper(BaseWrapper):
    
        @staticmethod
        def _time_method(method):
            name = "{0}.{1}.{2}".format(
                method.im_class.__module__,
                method.im_class.__name__,
                method.__name__
            )
    
            def timed_method(*args, **kwargs):
                begin = datetime.datetime.now()
                return_val = method(*args, **kwargs)
                end = datetime.datetime.now()
    
                trace = traceback.format_stack()
    
                db_log(
                    name,
                    begin,
                    end,
                    info={
                        'args': args,
                        'kwargs': kwargs,
                        'trace': trace
                    }
                )
                return return_val
    
            return timed_method
    
        def __init__(self, instance, methods):
            super(DBLogWrapper, self).__init__(instance)
            for method in methods:
                class_method = getattr(instance, method)
                wrapped_method = DBLogWrapper._time_method(class_method)
                setattr(self, method, wrapped_method)
    

    Output:

    >>> a = MyClass()
    >>> a = prova.DBLogWrapper(a, ['hello'])
    >>> a.hello()
    A
    ('__main__.MyClass.hello', datetime.datetime(2013, 1, 17, 20, 48, 26, 478023), datetime.datetime(2013, 1, 17, 20, 48, 26, 478071)) {'info': {'args': (), 'trace': ['  File "<stdin>", line 1, in <module>\n', '  File "prova.py", line 31, in timed_method\n    trace = traceback.format_stack()\n'], 'kwargs': {}}}
    >>> a.goodbye()
    B
    

    Anyway maybe you may use some __getattr__ magic instead, for example:

    class DBLogWrapper2(BaseWrapper):
    
        def __init__(self, instance, methods):
            super(DBLogWrapper, self).__init__(instance)
    
            self._methods = methods
    
        def __getattr__(self, attr):
            if attr not in methods:
                return getattr(self._instance, attr)
    
            def f(*args, **kwargs):
                return self.timed_method(getattr(self._item, attr),
                                         *args, **kwargs)
            return f
    
        def timed_method(method, *args, **kwargs):
            begin = datetime.datetime.now()
            return_val = method(*args, **kwargs)
            end = datetime.datetime.now()
    
            trace = traceback.format_stack()
    
            db_log(name,
                begin,
                end,
                info={
                    'args': args,
                    'kwargs': kwargs,
                    'trace': trace
                }
            )
            return return_val