Today we set out to implement a feature I saw and liked in Kotlin - Extension Methods.
You can follow along with working code samples here, or get the code here
Extension methods are a nice piece of syntactic-sugar that allow you to define free-functions and call them like instance methods. In Kotlin, it looks something like this:
Now, since they are free, static functions, they follow the same rules. They are not part of the class, nor have access to private members. And they can only be called in a scope where they are visible. Adding them in your code does not affect other code. Additionally, true member functions, if they exist, take precedence over extension methods (this is especially important with generic extension methods).
In our code today, we’ll try to mimic the features of extension methods as closely as possible. We’ll use the following syntax:
For extension methods, and the following implementation of
Square in our code throughout:
Monkey Patching 🙈
Python is a very dynamic language. Among other things, it allows us to change the attributes of (non-builtin) types at run-time. This means that we can extend our
Square class by adding a
draw method to it at run-time.
We’re now free to call
square.draw(). Before we discuss the draw-backs, let’s implement it with the syntax we defined:
Let’s go over this.
monkey_extend is a decorator with arguments. This is a common pattern where we use a decorator factory (
monkey_extend) to create a new decorator (
_decorator) as a closure, giving it access to the parameters passed to the factory (
cls). Then, in the core of the decorator, we use
setattr to do our monkey-patching.
While this works, it has several issues:
- Scope - once set, it can be used with any
Squarein any scope
- Precedence - it will override any existing
Dealing with precedence is easy (using
hasattr to check for existing
.draw) so we’ll focus on the scoping first.
Dynamic Attribute Lookup ✨
The first thing we know is that we need our new attribute to be there in some cases, and be gone in others - we need dynamic resolution. To do that, we’ll use
__getattr__. In Python classes,
__getattr__ is used in attribute lookup as a last resort, called when the other ways of looking up attributes came up empty. We’ll write our
__getattr__ along the following lines:
The first check,
has_extension, is basically checking whether the name we got matches the name of our extension method. Nothing to elaborate yet. Scoping, once again, remains the trickier part.
This is a bit much, so we’ll go over it in detail.
As a basis, we used the same decorator-with-parameters pattern here. We have
scoped_extend take the class we want to extend, then return
_decorator to get the job done. But instead of setting the attribute we want to extend, we monkey-patch
__getattr__ to our implementation (See (1)). This will override any existing implementation of
__getattr__, but we’ll get to that later. For now, we’ll focus on our implementation of
In (2) we implemented
has_extnesion - we simply compare the name we got to the name of our extension method. Then, in (3), comes some Python magic. Python allows us to inspect the running program, to see where we were called from and what variables were in scope in that code. To do that, we use the
inspect module. We use
inspect.stack() to get the call-stack for the current execution, then access the second frame (
) to get our caller. This will be where
getattr(obj, name) is invoked or
obj.name is used. We use
.frame to get the execution frame, and
f_globals to get the local and global variables available in that scope. They are equivalent to calling
locals() in the relevant frame.
With the scope at hand, we perform a lookup to see whether the extension method we defined is in that scope. To make sure we have our extension method, we get it by name, then ensure that it is truly our method.
Finally, in (4), when we know our method should be active, we bind it to the instance of the extended class and return it.
While our scope retrieval code works, it’s better to put it in a function rather than use it inline:
But, oh, we have to increment the stack index to
2 since we’re deeper in the callstack. This is risky. Instead, we’ll use the following trick to get the frame:
Instead of counting the frames in our code, changing them with every change - we’ll use the module system. We know that all of our scaffolding is in the same module, but the usage is not. This allows us to easily traverse the stack until we find code that does not belong in our module. That is our calling code.
Since you’re probably wondering - yes. You need to change
_get_first_external_stack_frame() if you want to put it in a different module. Implementing it is left as an exercise to the reader.
As mentioned before, our current implementation overrides any existing
__getattr__ function for the class. Lucky for us, fixing it is easy:
In (1) we get the original
__getattr__ method, to be stored for later usage. We use the
_default function to avoid an extra
if later. In (2) we use the saved
__getattr__, making sure that we only proceed to our code if it raised an
no_override_extend we have our first “to-spec” implementation of extension methods. We have both scoping and precedence down. It is time to celebrate and rest. But our quest is not done yet.
While our code works well for a proof-of-concept, there are still significant usability issues with it. Since the extension methods we create have nice and clean names, it is likely that we’ll want to use those names for other things. Unfortunately, once we do that, we’ll override the existing extension methods and they will no longer work:
The Fundemental Theorem of Software Engineering (FTSE) says that any problem can be solved by adding another level of indirection. Let’s see how this applies to our problem.
As mentioned in the interlude, our main issue is that of naming. Our extension method is bound to a name, and that name can be overriden in the scope that defines it. If that happens, we lose our extension method. To solve that, we’ll add another level of indirection - a scope that can safely hold our extension methods and protect them from being overriden. If you read our previous post you might recall that classes are wonderful for scopes. So we’ll use a class.
Our new syntax will look like this:
While we’re still using a decorator, you may notice that it takes no parameters. Instead, we use the extended type as the base type for our extension class. This allows us to write the extensions like any other subclass, with standard Python syntax, and then use the decorator to install the extensions in it.
Since we’ve already gone over the principles behind the construction of the decorator, let’s jump straight to the code and focus on the differences from the previous version:
First, you can see that there is no nested decorator - only the main one. And, as we mentioned before, we use inheritance to indicate which type we’re extending. So in (1) we access the base-class of our extension class to get the class we’re extending. Then, in (2) we check whether the requested attribute exists in our extension class. As you can see, the changes are pretty simple and straight-forward. In (3) we make the most important change - we check for the extension class in the scope, not the extension methods. This is the core of this change! And lastly, in (4), we get the required attribute from out extension class.
And with that, we’re done.
I hope you enjoyed this article. Regardless of that, I hope you never use it in production code.