Decorating Python

LEGO Decorator

Wrong decorator.

Ruby’s well known for its powerful metaprogramming features, and I recently talked about how I used them quite extensively to create smarter hashes. However, recently I’ve been working (playing?) in Python and I’ve learned that its implementation of something like a callback is much simpler. In Python, this is called a decorator, and if you’ve written much of it, you undoubtedly have already seen one in the wild – most commonly as @classmethod or @staticmethod.

However, writing your own decorators are powerful tools to have in your toolkit.

Let’s take a look at a simple implementation of a callback in both Ruby and Python and see how they compare.

For clarity, we’re looking to avoid having to do something like this in all our methods:

def method_with_callbacks
  do_before_callback

  # Actual method body

  do_after_callback
end

Ruby Callbacks

If you’ve worked with Rails, you should be familiar with ActiveSupport::Callbacks, or at least what they look like. You’re strongly encouraged to use them because they come pre-packaged with Rails itself. That’s great, unless you’re not using the framework and need to implement your callback in a plain Ruby script, like I did with intellihash.

To implement that functionality in plain Ruby, you need to create your callback in a module, then prepend it to the class you’re working with.

module Callbacks
  def before(method_name, callback)
    prepend(
      Module.new do
        define_method(method_name) do |*args, &block|
          send(callback)

          super(*args, &block)
        end
      end
    )
  end

  def after(method_name, callback)
    prepend(
      Module.new do
        define_method(method_name) do |*args, &block|
          result = super(*args, &block)
          
          send(callback)
  
          result
        end
      end
    )
  end
end

class ThingWithCallbacks
  extend Callbacks

  before :thing, :do_before_thing
  after :thing, :do_after_thing

  def thing
    puts 'Inside thing!'
  end

  private

  def do_before_thing
    puts 'Before thing!'
  end

  def do_after_thing
    puts 'After thing!'
  end
end

thing = ThingWithCallbacks.new
thing.thing
#=> Before thing!
#=> Inside thing!
#=> After thing!

Play with this code here.

Holy cow! That’s extraordinarily obtuse! We’re dynamically defining a module and a method, then prepending the method to the class to effectively overwrite the existing method in order to make this work. And it’s tricky on top of that, and requires you to understand how the class hierarchy of ThingWithCallbacks works:

ThingWithCallbacks.ancestors
#=> [#<Module:0x00007fb7e8976e48>, #<Module:0x00007fb7e8977050>, ThingWithCallbacks, Object, Kernel, BasicObject]

The first two modules are the dynamically defined modules we initiated to hold the methods we’ve re-defined. When Ruby is looking up the method we’re trying to use (in this case ThingWithCallbacks#thing) it first encounters it within these modules and never makes it to the original method definition. As long as we keep prepending modules, we’ll be able to effectively chain callbacks:

class ThingWithCallbacks
  extend Callbacks

  before :thing, :do_before_thing
  
  10.times { after :thing, :do_after_thing }

  # methods omitted...
end

thing = ThingWithCallbacks.new
thing.thing
#=> Before thing!
#=> Inside thing!
#=> After thing!
#=> After thing!
#=> After thing!
#=> After thing!
#=> After thing!
#=> After thing!
#=> After thing!
#=> After thing!
#=> After thing!
#=> After thing!

ThingWithCallbacks.ancestors
#=> [#<Module:0x00007ffbc52d6c98>, #<Module:0x00007ffbc52d6e00>, #<Module:0x00007ffbc52d7058>, #<Module:0x00007ffbc52d71c0>, #<Module:0x00007ffbc52d73a0>, #<Module:0x00007ffbc52d7580>, #<Module:0x00007ffbc52d7710>, #<Module:0x00007ffbc52d78f0>, #<Module:0x00007ffbc52d7ad0>, #<Module:0x00007ffbc52d7c38>, #<Module:0x00007ffbc52d7e40>, ThingWithCallbacks, Object, Kernel, BasicObject]

Wow. This is quite the mess. But it works!


Callbacks in Python with Decorators

We can easily implement the same thing in Python without messing with the inheritance chain by using the concept of decorators. These are special functions that return functions. They’re prepended by @ when you see them before a function definition, and they essentially wrap the function’s code and perform all kinds of operations. This is perfect for what we’re trying to do.

We’re going to abuse decorator chaining to create a “decorator generator” – here, decorators called before and after – that take a callback function as an argument and generate an appropriate secondary decorator to wrap the desired function in the callback. The decorator generator is here because passing an argument to a decorator means that it no longer automatically adds the function it wraps to the list of arguments, so is a good workaround to be able to essentially take both the function we want to wrap as well as a callback at the same time.

Let’s look:

class ThingWithCallbacks:
    def do_before_thing(fn):
        print('Before thing!')

    def do_after_thing(fn):
        print('After thing!')

    def before(callback):
        def before_decorator(fn):
            def wrapper(self):
                callback(self)
                fn(self)

            return wrapper
        return before_decorator

    def after(callback):
        def after_decorator(fn):
            def wrapper(self):
                fn(self)
                callback(self)

            return wrapper
        return after_decorator

    @before(do_before_thing)
    @after(do_after_thing)
    def thing(self):
        print('Inside thing!')


thing = ThingWithCallbacks()
thing.thing()
#=> Before thing!
#=> Inside thing!
#=> After thing!

Play with this code here.

The great thing here is you’ll notice we haven’t abused any of the class inheritance of ThingWithCallbacks. We’re simply hooking into the function definition in the class and adding before and after hooks to it.

You can also do exactly the same chaining as in Ruby…

class ThingWithCallbacks:
  # Lots of stuff omitted... 

  @before(do_before_thing)
  @after(do_after_thing)
  @after(do_after_thing)
  @after(do_after_thing)
  def thing(self):
      print('Inside thing!')

thing = ThingWithCallbacks()
thing.thing()
#=> Before thing!
#=> Inside thing!
#=> After thing!
#=> After thing!
#=> After thing!

Simplifying Things

Both of these snippets of code do identical things, but hopefully it can be seen that the Python version is somewhat simpler. It also has the advantage that it can be made even less complex if you’re willing to give up its genericness – i.e. if you only have one before callback, there’s no need to follow the decorator generator pattern, and you can just use a normal decorator instead.

Let’s take a look at a kata I was working on the other day. This solution asked to create an assembler in Python that would read in assembly instructions and output the correct values for the registers.

Computer hardware has the concept of a program counter, which tells it where in the instruction list it should currently read from. Manipulating the program counter is important as it’s the only way to proceed through the instructions. Normally, it’s incremented by 1 after every instruction – perfect to use a decorator:

class Assembler:
  def __init__(self, program):
    self.program = program
    self.registers = {}
    self.program_counter = 0
      
  def run(self):
    while self.program_counter < len(self.program):
      fn_name, *args = self.program[self.program_counter].split(" ")
      fn = getattr(self, fn_name)
      fn(*args)
          
  def increment_program_counter(fn):
    def wrapper(self, *args):
      fn(self, *args)
      self.program_counter += 1
        
    return wrapper
  
  @increment_program_counter
  def mov(self, register, value):
    self.registers[register] = self.register_to_value(value)
  
  @increment_program_counter
  def inc(self, register):
    self.registers[register] += 1
  
  @increment_program_counter
  def dec(self, register):
    self.registers[register] -= 1
      
  def jnz(self, test_value, inc_value):
    test_value = self.register_to_value(test_value)
    inc_value  = self.register_to_value(inc_value)
    
    if test_value != 0:
      self.program_counter += inc_value
    else:
      self.program_counter += 1

  # More helper instructions...

Play around with my toy assembler here.

Note that because the callback is exactly the same every single time, I don’t have to write the decorator generator, so my increment_program_counter decorator always does exactly the same thing after each function it wraps.

Note the functions mov, inc, dec being wrapped. Also note that the jnz function isn’t wrapped – this is Jump if Not Zero, which sets the program counter arbitrarily if the value being tested for isn’t zero. In this case, we can’t use the wrapper as that would interfere with the jump instruction.


Wrap-Up (Pun Intended)

I hope this was a neat example of how powerful Python decorators can be (along with module prepending in Ruby, to some extent). I plan on continuing to post neat metaprogramming tricks as I continue to discover them.

That’s all for now. Thanks for reading!