Search code examples
pythondecoratorpython-decorators

Decorating a Python class with a decorator as a class


Need some help to implement/understand how decorators as a class work in Python. Most examples I've found are either decorating a class, but implementend as a function, or implemented as a class, but decorating a function. My goal is to create decorators implemented as classes and decorate classes.

To be more specific, I want to create a @Logger decorator and use it in some of my classes. What this decorator would do is simply inject a self.logger attribute in the class, so everytime I decorate a class with @Logger I'll be able to self.logger.debug() in its methods.

Some initial questions:

  1. What does the decorator's __init__ receive as parameters? I it would receive only the decorated class and some eventual decorator parameters, and that's actually what happens for most of the cases, but please take a look at the output below for the DOMElementFeatureExtractor. Why does it received all those parameters?
  2. What about the __call__ method? What will it receive?
  3. How can I provide a parameter for the decorator (@Logger(x='y'))? Will it be passed to the __init__ method?
  4. Should I really be returning an instance of the class in the __call__ method? (only way I could make it work)
  5. What about chaining decorators? How would that work if the previous decorator already returned an instance of the class? What should I fix in the example below in order to be able to @Logger @Counter MyClass:?

Please take a look at this example code. I've created some dummy examples, but in the end you can see some code from my real project.

You can find the output at the end.

Any help to understand Python classes decorators implemented as a class would be much appreciated.

Thank you

from abc import ABC, abstractmethod

class ConsoleLogger:
  def __init__(self):
    pass
  
  def info(self, message):
    print(f'INFO {message}')

  def warning(self, message):
    print(f'WARNING {message}')
   
  def error(self, message):
    print(f'ERROR {message}')

  def debug(self, message):
    print(f'DEBUG {message}')

class Logger(object):
    """ Logger decorator, adds a 'logger' attribute to the class """
    def __init__(self, cls, *args, **kwargs):
      print(cls, *args, **kwargs)
      self.cls = cls
      
    def __call__(self, *args, **kwargs):
      print(self.cls.__name__)
      
      logger = ConsoleLogger()
      
      setattr(self.cls, 'logger', logger)
      
      return self.cls(*args, **kwargs)

class Counter(object):
    """ Counter decorator, counts how many times a class has been instantiated """
    count = 0
    def __init__(self, cls, *args, **kwargs):
       self.cls = cls
      
    def __call__(self, *args, **kwargs):
      count += 1
      
      print(f'Class {self.cls} has been initialized {count} times')
      
      return self.cls(*args, **kwargs)
      
@Logger
class A:
  """ Simple class, no inheritance, no arguments in the constructor """
  def __init__(self):
    self.logger.info('Class A __init__()')

class B:
  """ Parent class for B1 """
  def __init__(self):
    pass

@Logger
class B1(B):
  """ Child class, still no arguments in the constructor """
  def __init__(self):
    super().__init__()
    
    self.logger.info('Class B1 __init__()')
    
class C(ABC):
  """ Abstract class """
  def __init__(self):
    super().__init__()
    
  @abstractmethod
  def do_something(self):
    pass
  
@Logger
class C1(C):
  """ Concrete class, implements C """
  def __init__(self):
    self.logger.info('Class C1 __init__()')
  
  def do_something(self):
    self.logger.info('something')

@Logger
class D:
  """ Class receives parameter on intantiation """
  def __init__(self, color):
    self.color = color
    
    self.logger.info('Class D __init__()')
    self.logger.debug(f'color = {color}')

class AbstractGenerator(ABC):
  def __init__(self):
    super().__init__()
    
    self.items = None
    self.next_item = None
    
  @abstractmethod
  def __iter__(self):
    pass
  
  def __next__(self):
    pass
  
  def __len__(self):
    pass

  def __getitem__(self, key):
    pass
  
class AbstractDOMElementExtractor(AbstractGenerator):
  def __init__(self, parameters, content):
    super().__init__()
    
    self.parameters = parameters
    self.content = content
    
@Logger
class DOMElementExtractor(AbstractDOMElementExtractor):
  def __init__(self, parameters, content):
    super().__init__(parameters, content)
  
  def __iter__(self):
    self.logger.debug('__iter__')
  
  def __next__(self):
    self.logger.debug('__next__')  

  def __len__(self):
    self.logger.debug('__len__')

  def __getitem__(self, key):
    self.logger.debug('__getitem__')
    
class DOMElementFeatureExtractor(DOMElementExtractor):
  def __init__(self, parameters, content):
    super().__init__(parameters, content)

class DocumentProcessor:
  def __init__(self):
    self.dom_element_extractor = DOMElementExtractor(parameters={}, content='')
  
  def process(self):
    self.dom_element_extractor.__iter__()
    
a = A()
b1 = B1()
c1 = C1()
c1.do_something()
d = D(color='Blue')

document_processor = DocumentProcessor()
document_processor.process()

Output:

<class '__main__.A'>
<class '__main__.B1'>
<class '__main__.C1'>
<class '__main__.D'>
<class '__main__.DOMElementExtractor'>
DOMElementFeatureExtractor (<__main__.Logger object at 0x7fae27c26400>,) {'__module__': '__main__', '__qualname__': 'DOMElementFeatureExtractor', '__init__': <function DOMElementFeatureExtractor.__init__ at 0x7fae27c25840>, '__classcell__': <cell at 0x7fae27cf09d8: empty>}
A
INFO Class A __init__()
B1
INFO Class B1 __init__()
C1
INFO Class C1 __init__()
INFO something
D
INFO Class D __init__()
DEBUG color = Blue
DOMElementExtractor
DEBUG __iter__

Solution

  • Won't be a full answer, but I think it's helpful to review the basics of a decorator. This is what decorating looks like:

    @Logger
    class A:
      # A's code
    

    By definition, it's equivalent to doing this:

    class A
      # A's code
    
    A = Logger(A) # Logger has to be callable because...it's called
    

    Sources often say that decorators "modify", but that's really just the intended use. Technically, all you need is A to have a definition (so a function, method, or class) and Logger to be callable. If Logger returned "Hello, World", that's what A becomes.

    Okay, let's pretend we didn't decorate A for a bit and think about what it would take for Logger(A) to be "modifying." Well, A is a class, and you call a class to create instances: A(*args). Therefore, Logger(A)(*args) must also be instances of A. But Logger(A) isn't the class A, it's an instance of Logger. Luckily, you can make instances callable by defining the __call__ method in its class. Logger's __call__ method calls the class stored in its cls attribute and returns the instance.

    As for parameters in a decorator, it also helps to think about what it's equivalent to. You're interested in doing this:

    @Logger(x='y')
    class A:
      # A code
    

    So it's equivalent to this:

    class A:
      # A code
    
    A = Logger(x = 'y')(A)
    

    Note that Logger itself is not taking A as an argument. It's taking 'y' as an argument and returning another callable that takes A as an argument. So if Logger is a class, Logger(x = 'y') would be a Logger instance. Instances of a class can also serve as decorators if the class has a __call__ method!