Decorating Python
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!
- ← Previous: Metaprogramming Smarter Hashes in Ruby
- Next: Writing a Gameboy Game in 2021: Part 6 -- Postmortem →