I am trying to decorate a class

def decorate_module(Module):
    class Wrapper(Module):
        def __init__(self, cfg):
            self.c = create_c(**cfg)
            super().__init__()

    return Wrapper

While the code works, it gives a class named "Wrapper" instead of the name of the original module. When it comes to function decoration, one can use @wraps in this situation. I tried the same thing for class (by writing it just above class Wrapper(Module): but got an error.

What is the proper way to do this? A python built-in decorator @dataclass seems to be handling this well. I wonder how it works.. Is there a helper function like @wraps? Or should I manually overwrite __name__? If so, what are the full list of such "double-underscore "properties?

(EDIT) my original purpose

Given a class that does not take any arguments in __init__, I would like to decorate it to take an argument "cfg", do something (create_c), and store it as an attribute.

To be precise, my decorator will actually take "create_c" as an argument (I omitted this in the above code since additional nesting would make it verbose). Then, I am thinking of using it as follows.

@decorate_module(create_fn1)
class Module1:
    ...

@decorate_module(create_fn2)
class Module2:
    ...

CodePudding user response：

There is no built-in equivalent to functools.wraps for a class, but you can copy yourself the dunder attributes of the class being wrapped, which include __doc__, __name__, __qualname__ and __module__, as documented.

def decorate_module(Module):
    class Wrapper(Module):
        def __init__(self, cfg):
            self.c = create_c(**cfg)
            super().__init__()

    for attr in '__doc__', '__name__', '__qualname__', '__module__':
        setattr(Wrapper, attr, getattr(Module, attr))
    return Wrapper

CodePudding user response：

As far as I understood you need kind of "subclasser" so you don't need at all the parameter module (see my comment). I used __new__ to mimic the decoration style, it dynamically creates a subclass with the type build-in function. The __init__ is more painful: it requires an extra function because __init__ must always return None

class my_dataclass:
    def __new__(cls, tcls):
        sub_cls = type(f'{tcls.__name__}Jr', (tcls,), {})

        def constructor(self, **cfg):
            setattr(self, 'c', super(type(self), self).create_c(**cfg))
            super(type(self), self).__init__()
            
        setattr(sub_cls, '__init__', lambda self, **cfg: constructor(self, **cfg))
        
        return sub_cls


@my_dataclass
class A:
    
    def __init__(self):
        print(self, '__init__')

    def create_c(self, **cfg):
        print(**cfg)
        return 'create_c'
        
print(A)
a = A()
print(a.c) # attribute

Output

<class '__main__.AJr'>

<__main__.AJr object at 0x7f4105dcc2e0> __init__
create_c