Search code examples
pythondecoratorpython-dataclasses

How to use `@dataclass` decorator within a decorator?


I would like to create a class decorator to pre-define some attributes and methods to decorated class. At the same time, I would like that this decorated class be decorated with @dataclass decorator. To make things even easier for the user, I would like the user to have only to use my decorator so that his/her class is decorated with @dataclass.

So here is what I have come up with to decorate a class, so that the decorated class has a new attribute id_. I am submitting this code to know if the nesting of decoartor is 'acceptable'? Basically, I am not so sure about calling a decorator like a function (while it is a function), so I wonder if there could be side effects?

from dataclasses import dataclass

# My library: complexitiy is allowed.
def indexer(orig_class):
    orig_class = dataclass(orig_class)
    # Copy of original __init__ to call it without recursion.
    orig_init = orig_class.__init__
    id_ = 10

    def __init__(self,*args, **kws):
        self.id_ = id_
        orig_init(self, *args, **kws) # Call the original __init__

    orig_class.__init__ = __init__ # Set the class' __init__ to the new one

    return orig_class

# User environment: importing @indexer decorator: ease of use is targeted.
# Expected behavior is that @indexer provides same functions as @dataclass.
@indexer
class Truc:
    mu : int

# Test
test = Truc(3)

test.mu
Out[37]: 3

test.id_
Out[38]: 10

As a sidenote, I don't know if decorator is better than class inheritance here, when it comes to adding attributes and methods? Thanks for the advices! Bests,

Edit

To whom may read till here, and wonder about use between decorator and inheritance.

Going farther in the implementation, I finally stick with decorator as it seems to me it provides a more flexible way of tweaking what would have been a child class if had retained inheritance.

With inheritance, I don't see how the parent class has any control on child attributes (to make them frozen for instance) while leaving complete freedom on how naming/defining these attributes.

But this is actually an important property I am looking for with dataclass.

So, current status is reviewable with below @indexerdecorator. It leaves freedom to the user to create a 'data class' with the number of parameters he/she wants, naming them as he/she wants. But it makes possible

  • to compare and check equality between 2 instances of decorated class.
  • once an instance is created from decorated class, it cannot be modified (frozen)
  • it equips the decorated class with a (modifiable when decorating) attribute fields_sep
  • it equips the decorated class with a method to print as string parameters value defined when instancing the decorated class, and not those added by the decorator (like fields_sep)

It fulfills my requirements (I can use the objects created from decorated class) as keys for a dict or sorted dict.

In case it can help, here it is.

PS: if you have any advises or better practices, I welcome them.

from dataclasses import dataclass, asdict

# for optional arguments in decorator: https://stackoverflow.com/a/24617244/4442753
def indexer(index_class=None, *, fields_sep:str='.'):
    def _tweak(index_class):
        # Wrap with `@dataclass`.
        index_class = dataclass(index_class, order= True, frozen=True)
        # Copy of original __init__ to call it without recursion.
        index_class_init = index_class.__init__
    
        def __init__(self,*args, **kws):
            object.__setattr__(self, "_fields_sep", fields_sep)
            index_class_init(self, *args, **kws)
    
        index_class.__init__ = __init__
    
        def _to_str(self):
            return self._fields_sep.join(map(str, asdict(self).values()))
    
        index_class._to_str = property(_to_str)
    
        return index_class
    if index_class:
        # Calling decorator without other parameters.
        return _tweak(index_class)
    # Calling decorator with other parameters.
    return _tweak

# Test without parameter.
@indexer
class Test:
    mu : int
    nu : str
test = Test(3, 'oh')
assert test._to_str == '3.oh'

# Test with 'fields_sep' parameter.
@indexer(fields_sep='-')
class Test:
    mu : int
    nu : str
test = Test(3, 'oh')
assert test._to_str == '3-oh'

# Test (un)equality.
@indexer
class Test:
    mu : int
    nu : str
test1 = Test(3, 'oh')
test2 = Test(3, 'oh')
assert test1 == test2
test3 = Test(4, 'oh')
assert test1 != test3

# Test dict.
di = {test1:[1,2,3], test2:[7,8,9]}
assert len(di) == 1
assert di[Test(3, 'oh')] == [7,8,9]

Solution

  • Basically, I am not so sure about calling a decorator like a function (while it is a function), so I wonder if there could be side effects?

    decorators in python are syntactic sugar, PEP 318 in Motivation gives following example

    def foo(cls):
        pass
    foo = synchronized(lock)(foo)
    foo = classmethod(foo)
    

    is equivalent to

    @classmethod
    @synchronized(lock)
    def foo(cls):
        pass
    

    In other word decorators allow you to write less lines of codes for getting very same result.