TIL: Writing decorators for classes

To my surprise writing decorators for classes is easier than for functions. Here's how to do it in annotated fashion with an unnecessary decorator that doesn't accept any additional arguments.

# Write a callable that accepts a cls as an argument 
def tools(cls):
    # Write functions that accept "self: object" as an argument.
    def simplistic_attribute_count(self: object) -> int:
        """Returns the number of attributes."""
        return len(self.__dict__)

    def docs(self: object) -> str:
        """Returns the docstring for the class."""
        return self.__doc__

    # Attach the functions as methods
    cls.simplistic_attribute_count = simplistic_attribute_count
    cls.docs = docs

    # Return the modified class
    return cls

Let's test it out:

@tools
class A:
    """Docstring for testing the tools decorator"""


a = A()
a.one = 1

assert a.simplistic_attribute_count() == 1
assert a.docs() == 'Docstring for testing the tools decorator'

Next up, how to do this while passing in arguments!