How to use the Class decorator wrapper?

Viewed 97

I have a data_class object and want to add some functionality using a decorator. It works, however the class name is lost and becomes defaultinstance.<locals>.SubClass .

I know functools.wraps can fix this, but where/how should it be used?

from dataclasses import dataclass, fields, is_dataclass
   
def defaultinstance(object):
  
    class SubClass(object):
        @classmethod
        def woof(cls):
            print("woof")

    return SubClass


@defaultinstance
@dataclass
class Dog:
    name: str
    paws: int


doggo = Dog(name="jack", paws=4)

print(doggo)
# defaultinstance.<locals>.SubClass(name='jack', paws=4)
1 Answers

You can just have the decorator modify the class like this

from dataclasses import dataclass, fields, is_dataclass
   
def defaultinstance(object):
    """
    creates a class method 'woof' and adds it to object
    """
    @classmethod
    def woof(cls):
        print("woof")
    
    setattr(object, 'woof', woof)
    return object


@defaultinstance
@dataclass
class Dog:
    name: str
    paws: int


doggo = Dog(name="jack", paws=4)

print(doggo)
Dog.woof()  # as it's a class method we can use it directly on Dog
doggo.woof()  # but it works fine on doggo too

with result

Dog(name='jack', paws=4)
woof
woof

This way you don't have to create a new class that would need to masquerade as the old class with the old class never being used.

Related