pythonpython-typingpython-decoratorspython-dataclasses

Re-decorate a python (class) decorator


I'd like to create a decorator that basically wraps an already existing decorator that has parameters, such that the new decorator acts like the old one with some of the arguments supplied.

Specifically, this is about the builtin @dataclass decorator. I have a number of classes to decorate with it, while always using kw_only=True and eq=False, and I'd like to have a new decorator that does just that, saving me to spell out the parameters every time. So, e.g. ...

@mydataclass
class Foo:
    a: int = 5

...should be equivalent to...

@dataclass(kw_only=True, eq=False)
class Foo:
    a: int = 5

I know this doesn't seem like saving a lot of typing, but well this is more about providing a convenience decorator for the rest of our team, so that no one forgets to correctly add these two parameters.


Solution

  • @dataclass is handled specially by type checkers. The main reason is to infer the __init__ signature, but also other behavior. This is not inferred by normal decorators that wrap dataclass.

    For correct handling by type checkers you need to decorate your decorator with typing.dataclass_transform which tells type checkers that your decorator is like @dataclass

    from dataclasses import dataclass, field
    
    @typing.dataclass_transform(
        field_specifiers=(field,)
        kw_only_default=True,
        eq_default=False,
    )
    def mydataclass(cls):
        return dataclass(kw_only=True, eq=False)(cls)
    
    @mydataclass
    class Foo:
        a: int = 5
    
    Foo(a=1)  # ok