How to use Python function decorators in a class?

function_decorator_in_a_class.py

class Lightning:
  lights_on = 'ON'
  lights_off = 'OFF'

  def button_action(action_type):
    def action_decorator(func):
      def action_wrapper(self):
        if action_type == 'on':
          print(self.lights_on)
        func(self)
        if action_type == 'off':
          print(self.lights_off)
      return action_wrapper
    return action_decorator

  @button_action('off')
  def enlight_off(self):
    print('Enlight me!')

  @button_action('off')
  def enlight_on(self):
    print('Enlight me!')

if __name__ == '__main__':
  test = Lightning()
  test.enlight_on()
  print('-----')
  test.enlight_off()