Search code examples
pythondecoratorpython-decorators

Python: How can i decorate function to change it into class method


I have code like this and i want to write decorator which will add decoradted function as class method of class A.

class A:
    pass

@add_class_method(A)
def foo():
    return "Hello!"

@add_instance_method(A)
def bar():
    return "Hello again!"

assert A.foo() == "Hello!"
assert A().bar() == "Hello again!"

Solution

  • What about this approach?
    P.S. The code is not structurally optimized for the sake of clarity

    from functools import wraps
    
    
    class A:
        pass
    
    
    def add_class_method(cls):
        def decorator(f):
            @wraps(f)
            def inner(_, *args, **kwargs):
                return f(*args, **kwargs)
    
            setattr(cls, inner.__name__, classmethod(inner))
    
            return f
    
        return decorator
    
    
    def add_instance_method(cls):
        def decorator(f):
            @wraps(f)
            def inner(_, *args, **kwargs):
                return f(*args, **kwargs)
    
            setattr(cls, inner.__name__, inner)
    
            return f
    
        return decorator
    
    
    @add_class_method(A)
    def foo():
        return "Hello!"
    
    
    @add_instance_method(A)
    def bar():
        return "Hello again!"
    
    
    assert A.foo() == "Hello!"
    assert A().bar() == "Hello again!"