Add a decorator to a class add a decorator to class' methods by decorating class

Question

I am trying to create a decorator which can be defined on the class and decorates everything defined in it. First let me show the setup that I got already based on other SO answers:

import inspect


# https://stackoverflow.com/a/18421294/577669
def log(func):
    def wrapped(*args, **kwargs):
        try:
            print("Entering: [%s]" % func)
            try:
                # https://stackoverflow.com/questions/19227724/check-if-a-function-uses-classmethod
                if inspect.ismethod(func) and func.__self__:  # class method
                    return func(*args[1:], **kwargs)
                if inspect.isdatadescriptor(func):
                    return func.fget(args[0])
                return func(*args, **kwargs)
            except Exception as e:
                print('Exception in %s : (%s) %s' % (func, e.__class__.__name__, e))
        finally:
            print("Exiting: [%s]" % func)
    return wrapped


class trace(object):
    def __call__(self, cls):  # instance, owner):
        for name, m in inspect.getmembers(cls, lambda x: inspect.ismethod(x) or inspect.isfunction(x)):
            setattr(cls, name, log(m))
        for name, m in inspect.getmembers(cls, lambda x: inspect.isdatadescriptor(x)):
            setattr(cls, name, property(log(m)))
        return cls


@trace()
class Test:
    def __init__(self, arg):
        self.arg = arg

    @staticmethod
    def static_method(arg):
        return f'static: {arg}'

    @classmethod
    def class_method(cls, arg):
        return f'class: {arg}'

    @property
    def myprop(self):
        return 'myprop'

    def normal(self, arg):
        return f'normal: {arg}'


if __name__ == '__main__':
    test = Test(1)
    print(test.arg)
    print(test.static_method(2))
    print(test.class_method(3))
    print(test.myprop)
    print(test.normal(4))

When removing the @trace decorator from the class, this is the output:

123
static
class
myprop
normal

When adding the @trace decorator I get this:

Entering: [<function Test.__init__ at 0x00000170FA9ED558>]
Exiting: [<function Test.__init__ at 0x00000170FA9ED558>]
1
Entering: [<function Test.static_method at 0x00000170FB308288>]
Exception in <function Test.static_method at 0x00000170FB308288> : (TypeError) static_method() takes 1 positional argument but 2 were given
Exiting: [<function Test.static_method at 0x00000170FB308288>]
None
Entering: [<bound method Test.class_method of <class '__main__.Test'>>]
Exiting: [<bound method Test.class_method of <class '__main__.Test'>>]
class: 3
Entering: [<property object at 0x00000170FB303E08>]
Exiting: [<property object at 0x00000170FB303E08>]
myprop
Entering: [<function Test.normal at 0x00000170FB308438>]
Exiting: [<function Test.normal at 0x00000170FB308438>]
normal: 4

Conclusion from this example: the init, normal, class and prop methods are all instrumented correctly.

However, the static method is not.

My questions for this snippet are:

Is it ok to check certain usecases like I did in the log? Or is there a better way?
How to see if something is a static method to be able to pass in nothing (because now the Test-instance is passed in)?

Thanks!

Testing my answer, and I notice your output doesn't match what I'd expect in your code. Did you change some test data? — Kenny Ostrom, Dec 31 '21 at 15:28

Kenny Ostrom · Answer 1 · 2021-12-31T16:27:08.697

I looked at the source of inspect, and noticed that it can find static methods in classify_class_attrs, so I modified your code to use that function.

I also separated the log so I could have different wrapper functions to handle the different rules. Some of them are redundant, but that's how I separated staticmethod, originally. I was worried that classmethod should get the cls argument, and maybe that's a legitimate concern, but it passed these simple tests without that becoming an issue.

import inspect
import types

# https://stackoverflow.com/a/18421294/577669
def log(func, *args, **kwargs):
    try:
        print("Entering: [%s]" % func)
        try:
            if callable(func):
                return func(*args, **kwargs)
        except Exception as e:
            print('Exception in %s : (%s) %s' % (func, e.__class__.__name__, e))
            raise e
    finally:
        print("Exiting: [%s]" % func)

def log_function(func):
    def wrapped(*args, **kwargs):
        return log(func, *args, **kwargs)
    return wrapped
    
def log_staticmethod(func):
    def wrapped(*args, **kwargs):
        return log(func, *args[1:], **kwargs)
    return wrapped
    
def log_method(func):
    def wrapped(*args, **kwargs):
        instance = args[0]
        return log(func, *args, **kwargs)
    return wrapped
    
def log_classmethod(func):
    def wrapped(*args, **kwargs):
        return log(func, *args[1:], **kwargs)
    return wrapped

def log_datadescriptor(name, getter):
    def wrapped(*args, **kwargs):
        instance = args[0]
        return log(getter.fget, instance)
    return wrapped
    
class trace(object):
    def __call__(self, cls):  # instance, owner):
        for result in inspect.classify_class_attrs(cls):
            if result.defining_class == cls:
                func = getattr(cls, result.name, None)
                if result.kind == 'method':
                    setattr(cls, result.name, log_method(func))
                if result.kind == 'class method':
                    setattr(cls, result.name, log_classmethod(func))
                if result.kind == 'static method':
                    setattr(cls, result.name, log_staticmethod(func))
        for name, getter in inspect.getmembers(cls, inspect.isdatadescriptor):
            setattr(cls, name, property(log_datadescriptor(name, getter)))
        return cls


@trace()
class Test:
    def __init__(self, arg):
        self.value = arg

    @staticmethod
    def static_method(arg):
        return f'static: {arg}'

    @classmethod
    def class_method(cls, arg):
        return f'class Test, argument: {arg}'

    @property
    def myprop(self):
        return f'myprop on instance {self.value}'

    def normal(self, arg):
        return f'normal: {arg} on instance {self.value}'


if __name__ == '__main__':
    test = Test(123)
    print(test.value)
    print(test.static_method(2))
    print(test.class_method(3))
    print(test.myprop)
    print(test.normal(4))

Entering: [<function Test.init at 0x000002338FDCCA60>]
Exiting: [<function Test.init at 0x000002338FDCCA60>]
123
Entering: [<function Test.static_method at 0x000002338FDCCAF0>]
Exiting: [<function Test.static_method at 0x000002338FDCCAF0>]
static: 2
Entering: [<bound method Test.class_method of <class 'main.Test'>>]
Exiting: [<bound method Test.class_method of <class 'main.Test'>>]
class Test, argument: 3
Entering: [<function Test.myprop at 0x000002338FDCCC10>]
Exiting: [<function Test.myprop at 0x000002338FDCCC10>]
myprop on instance 123
Entering: [<function Test.normal at 0x000002338FDCCCA0>]
Exiting: [<function Test.normal at 0x000002338FDCCCA0>]
normal: 4 on instance 123

Some of the text doesn't match exactly because we both made some trivial changes to the output in the traced class.

I noticed that this also traces __weakref__ Maybe it should skip any dunder names? — Kenny Ostrom, Dec 31 '21 at 15:53
Also, I wasn't sure how to handle the property getter from classify_class_attrs's result, so I just copied your already working code for handling those. — Kenny Ostrom, Dec 31 '21 at 16:23
Thanks for the method "classify_class_attrs". One issue I'm having is when decorating both base and derived classes. Do you see a way around that? — Steven Van Ingelgem, Jan 01 '22 at 04:39
I'll look into that if you're satisfied with the staticmethod answer. — Kenny Ostrom, Jan 01 '22 at 15:08
thanks a lot for the hints. I think I figured it out in a reasonable clean way. I posted mine as a solution, but am unsure what to do. Your hint for the classify_class_attrs did bring me to the solution. Please do tell me how to proceed. — Steven Van Ingelgem, Jan 01 '22 at 19:59

score 0 · Answer 2 · answered Jan 01 '22 at 19:58

My final solution:

import inspect
from typing import Type

from decorator import decorator


@decorator
def log(func, *args, **kwargs):
    try:
        print("Entering: [%s]" % func)
        return func(*args, **kwargs)
    finally:
        print("Exiting: [%s]" % func)


def _apply_logger_to_class(cls):
    for attr in inspect.classify_class_attrs(cls):
        if attr.defining_class is not cls:
            continue
        if attr.kind == 'data':
            continue

        if isinstance(attr.object, (classmethod, staticmethod)):
            setattr(cls, attr.name, attr.object.__class__(log(attr.object.__func__)))
        elif isinstance(attr.object, property):
            setattr(cls, attr.name, property(log(attr.object.fget)))
        else:
            setattr(cls, attr.name, log(attr.object))

    return cls


def trace(func=None):
    if isinstance(func, type):
        return _apply_logger_to_class(func)  # logger is the class

    return log(func)


@trace
def normal_function_call(arg):
    return f'normal_function_call: {arg}'


@trace
class Test:
    def __init__(self, arg):
        self.arg = arg
        print(f'{self.__class__.__name__}.__init__: {arg}')

    @staticmethod
    def static_method(arg):
        return f'Test.static: {1}'

    @classmethod
    def class_method(cls, arg):
        print(f'{cls.__name__}.class: {arg}')

    @property
    def myprop(self):
        print(f'{self.__class__.__name__}.myprop.getter')
        return 1

    @myprop.setter
    def myprop(self, item):
        print(f'{self.__class__.__name__}.myprop.setter')

    @myprop.deleter
    def myprop(self):
        print(f'{self.__class__.__name__}.myprop.deleter')

    def normal(self, arg):
        print(f'{self.__class__.__name__}.normal: {arg}')


@trace
class TestDerived(Test):
    @staticmethod
    def static_method(arg):
        print(f'TestDerived.class: {arg}')


if __name__ == '__main__':
    print(normal_function_call(0))

    def do_test(test_class: Type[Test]):
        print('-'*20, test_class.__name__)

        test = test_class(1)
        test.static_method(2)
        test.__class__.static_method(2.5)
        test.class_method(3)
        test.__class__.class_method(3.5)
        test.myprop
        test.normal(4)

        assert inspect.getfullargspec(test.normal).args == ['self', 'arg']
        assert inspect.getfullargspec(test.normal).kwonlyargs == []

    do_test(Test)
    do_test(TestDerived)

This works as intended for derived classes, for all objects I want, and preserves the signatures. (which @wraps does not).

Add a decorator to a class add a decorator to class' methods by decorating class

2 Answers2