Writing/Metaprogramming: An Adventure
Chapter 3 · 23 min read
Bytecode Manipulation
Let's pick another challenge that may seem useless, and see what lessons it might teach us. One day, I wrote the following class:
>>> class A:
... def __init__(self):
... self.x = 1
... def f(self):
... return self.x
And got annoyed: why do I have to include self in every method signature?
Languages like C++ and Java have an implicit this — and yet my beautiful
Python doesn't?
Selfless Classes
So I wrote this instead:
>>> class A:
... def __init__():
... self.x = 1
... def f():
... return self.x
And set out on a quest to make it work, calling forth the powers of decorators,
descriptors, even tracers. But to no avail: for the life of me, I couldn't
change the local self variable to automatically locate the instance — because,
when I looked at the methods' bytecode, I realized that there was no local
variable self to begin with.
You see, the moment I omitted it from the signature, Python decided it must be a global variable, and understandably so; at no point was it defined at or passed into the local scope:
>>> import dis
>>> dis.dis(A)
Disassembly of __init__:
5 0 RESUME 0
6 2 LOAD_CONST 1 (1)
4 LOAD_GLOBAL 0 (self)
14 STORE_ATTR 1 (x)
24 RETURN_CONST 0 (None)
Disassembly of f:
8 0 RESUME 0
9 2 LOAD_GLOBAL 0 (self)
12 LOAD_ATTR 2 (x)
32 RETURN_VALUE
Interlude: Code Objects
To fix it, then, we'd have to rewrite the methods' bytecode; and to do that, a
brief overview of how code objects are organized is in order. When we write a
function, even before we call it, Python goes ahead and translates it into a
series of smaller steps, which the interpreter can execute one by one. For
example, the line self.x = 1 is translated into 3 such steps:
-
LOAD_CONST, which brings the constant (in this case,1) into focus; -
LOAD_GLOBAL, which brings a global name (in this case,self) into focus; -
And
STORE_ATTR, which takes an attribute name (in this case,x), and sets that attribute of the name-in-focus to the value-in-focus (that is,self.x = 1).
The reason for this verbosity is that opcodes are stored in dense, binary form, and carried out by a lower-level implementation (in case of CPython, it's in C). Since it can't handle complex, high-level abstractions, those need to be broken down into simpler and more generic operations, much like compilation of C code into an executable file requires generating machine instructions that the hardware can actually process. On one hand, this introduces a complicated translation process; on the other, it means that the same code can run in different environments, as long as we have an appropriate Python interpreter (or C compiler) to do the necessary adaptations for each.
As for the "focus" I mentioned: it's simply a way to support opcodes with
multiple arguments. You see, each instruction is encoded in just two bytes —
one to denote the operation (like LOAD_CONST, which happens to be 100, or
LOAD_GLOBAL, which happens to be 116), and one to denote its argument (like
1 or self). This means that anything that requires a broader context (like
STORE_ATTR/95) will assume that previous operations have prepared the
necessary data for it, and blindly consume it from an agreed-upon "drop point".
This drop point is effectively a stack: LOAD_CONST and LOAD_GLOBAL push a
value each unto it, and those are exactly the two values STORE_ATTR pops off,
so everything checks out. Of course, none of this matters if we use Python as
intended; bytecode is compiled automatically, behind the scenes, in which case
these conventions are guaranteed to be respected. It's only when we decide to
work at this level ourselves that such details start to matter, because unless
we coordinate everything very carefully, some truly bizarre results might ensue.
One final note (for now): since the argument must fit in a single byte, it can't
possibly store the entire string value "self"; even the number 1 would only
fit because it's so small — with self.x = 1000 we'd have the same problem. For
that reason, the opcode argument is never the value itself, but rather an index
into a separate collection, determined by the operation in question. That's why
dis shows LOAD_CONST 1 and LOAD_GLOBAL 0, with the dereferenced values
added in parentheses.
To make sure we understand this, let's try to parse everything ourselves:
>>> import opcode # maps operation codes to names
>>> code = A.__init__.__code__
>>> code.co_code # that's the bytecode:
b'\x97\x00d\x01t\x00\x00\x00\x00\x00\x00\x00\x00\x00_\x01...'
>>> code.co_code[0], code.co_code[1] # the 1st instruction
(151, 0)
>>> opcode.opname[151] # let's skip this for now…
'RESUME'
>>> code.co_code[2], code.co_code[3] # the 2nd instruction
(100, 1)
>>> opcode.opname[100] # a constant…
'LOAD_CONST'
>>> code.co_consts[1] # … so its value is here
1
>>> code.co_code[4], code.co_code[5] # the 3rd instruction
(116, 0)
>>> opcode.opname[116] # a global…
'LOAD_GLOBAL'
>>> code.co_names[0] # … so its value is here
'self'
>>> code.co_code[6], code.co_code[7] # let's skip all the zeros for now…
(0, 0)
>>> code.co_code[14], code.co_code[15] # the 4th instruction
(95, 1)
>>> opcode.opname[95] # attribute assignment
'STORE_ATTR'
>>> code.co_names[1] # … so its name is here
'x'
As you can see, it's pretty tedious; luckily, it's all part of the Python
standard, so we have builtin utilities to make our life easier. We've already
seen dis, which stands for disassembly — reversing bytecode into a more
readable format. But like with any language, understanding is considerably
easier than speaking ourselves; so before we try to rewrite it, let's address
the nuances that we skipped.
Bytecode Optimizations
One thing that's important to take into account when working with bytecode is that it's somewhat fragile: there's no guarantee that it will stay the same between major Python versions, so parsing it manually can easily break. Even in our exceedingly simple case, compare the previous disassembly, with Python 3.12, to this one with Python 3.10:
>>> dis.dis(A.__init__)
Disassembly of __init__:
7 0 LOAD_CONST 1 (1)
2 LOAD_GLOBAL 0 (self)
4 STORE_ATTR 1 (x)
6 LOAD_CONST 0 (None)
8 RETURN_VALUE
Several things stand out: there's no RESUME instruction, and rather than
RETURN_CONST (a new, optimized opcode), we have the older RETURN_VALUE,
which uses the focus paradigm (a.k.a. "leave me the arguments on the stack"), so
it's preceded by a LOAD_CONST for that purpose (even though in our case it's
an implicit return statement, so its value is simply None).
A subtler difference is the padding: the number next to each operation is its
offset in the bytecode. Since we said that each instruction is exactly 2 bytes
long, those offsets are 0, 2, 4, 6 and 8, as we might expect; and
looking at the bytecode itself:
>>> A.__init__.__code__.co_code # python 3.10
b'd\x01t\x00_\x01d\x00S\x00'
We see it's missing all those zeroes we had before. In Python 3.12 we jumped
from 4 LOAD_GLOBAL to 14 STORE_ATTR, then 24 RETURN_CONST – why is that?
These are all optimizations. Bytecode is not intended to be simple or convenient — it's much more important that it runs smoothly and efficiently, which means that we have to tolerate any quirks that might arise as a result. Even if they make little sense in our case, we have to accept these complexities and accommodate them.
Take the RESUME instruction, for example: it doesn't matter for a function
like ours, and didn't exist prior to Python 3.11. But in the case of generators
or coroutines, which can yield or await at any point, only to resume
execution later — this is exactly what helps the interpreter orient itself
faster. As for the zeroes, known as CACHE instructions — they are placeholders
waiting to be overwritten, if at any point Python detects a pattern in our code
and decides to dynamically optimize it. Without them, it would have to take the
entire bytecode apart and rearrange it, fixing any offsets that might have
shifted as a result; with them, it essentially has some breathing room — a
buffer into which to squeeze its changes without having to move everything else
around.
All of this doesn't really matter for our purposes, but as we shall see shortly, it's important to understand and respect; working at such a low-level has its advantages, but also comes at a price.
Rewriting Bytecode
Now that we understand everything that happens within our method, we can come up
with a plan on how to change it. The first and most obvious point would be
replacing any LOAD_GLOBAL that references self with LOAD_FAST (that's the
instruction to load a local name) — but with what argument?
That's the second point, then: we'd have to add self to the co_varnames
collection used by LOAD_FAST; and more than that, we'd have to prepend it to
the very beginning, since Python methods work by implicitly injecting self as
the first argument. This, in turn, means that if we have any other local
arguments — which would be referenced by other LOAD_FAST instructions in our
bytecode — we'd have to bump their argument's index by one to account for this
shift. In our case, it doesn't matter; but a small change in the code:
>>> class A:
... def __init__(x):
... self.x = x
Would show you what I mean:
>>> dis.dis(A.__init__) # python 3.10
Disassembly of __init__:
7 0 LOAD_FAST 0 (x)
2 LOAD_GLOBAL 0 (self)
4 STORE_ATTR 1 (x)
6 LOAD_CONST 0 (None)
8 RETURN_VALUE
If self is to be our first local, at index 0, then it would push x to
index 1, and we must fix this instruction to be LOAD_FAST 1.
To do all that, we'll use get_instructions: it's similar to dis in that it
does all the disassembly for us, but instead of printing it to the screen, we
get convenient Instruction objects that we can work with. So:
>>> code = A.__init__.__code__
>>> new_varnames = list(code.co_varnames)
>>> new_varnames.insert(0, "self")
>>> new_bytecode = []
>>> for i in dis.get_instructions(code):
... op, arg = i.opcode, i.arg
... if i.opname == "LOAD_GLOBAL" and i.argval == "self":
... op, arg = opcode.opmap["LOAD_FAST"], 0
... elif i.opname == "LOAD_FAST":
... arg += 1
... new_bytecode.extend([op, arg or 0]) # in case arg is None
All that's left to do is update __init__; but since code objects are
immutable, we can't just patch in our new values. Creating a brand new code
object is also a bit annoying, since there are almost 20 arguments we'd have to
pass in to configure it properly — but conveniently, existing code objects have
a replace method (at least, since Python 3.8), which takes only those
configurations we want to change, and copies over the rest. Besides the new
bytecode and local names, we also need to update a couple more fields:
co_nlocals, which indicates the number of local variables, and co_argcount,
which indicates how many of those are to be considered arguments. And so:
>>> A.__init__.__code__ = code.replace(
... co_code=bytes(new_bytecode),
... co_varnames=tuple(new_varnames),
... co_nlocals=code.co_nlocals + 1,
... co_argcount=code.co_argcount + 1,
... )
Drum roll:
>>> a = A(1)
segmentation fault
Rewriting Bytecode Optimizations
That was a bit underwhelming; although we can take comfort in that it's not so easy to get a segmentation fault in Python, so we definitely touched a nerve. If we would switch our version to 3.10, this would actually work:
>>> a = A(1) # python 3.10
>>> a.x
1
But in Python 3.12, the fact that we ignored the CACHE padding (effectively
removing it from our new bytecode) made the interpreter literally crash. Like I
said: those optimizations, annoying though they are, must be respected.
We can try to preserve the padding by filling in gaps based on each instruction's offset:
>>> new_bytecode = []
>>> for i in dis.get_instructions(code):
... if len(new_bytecode) < i.offset:
... new_bytecode.extend([0] * (i.offset - len(new_bytecode)))
... # … same as before
But that wouldn't work either, because the padding actually depends on the particular instruction preceding it — and since we're replacing them on the fly, we'd have to adjust the padding judiciously, not just copy it over as-is.
Now, what's challenging about this situation is that there isn't much documentation on how exactly any of this works — one has to sieve through PEPs and discussion groups, or read the interpreter's source code outright. The simplest approach, though, which I'll show here, is naive reverse-engineering: Python being the dynamic language that it is, we can just try some things, derive a pattern, apply it and see if it works. We know the code we want to end up with, so let's just look at its disassembly:
>>> class A:
... def __init__(self, x):
... self.x = x
>>> dis.dis(A.__init__)
Disassembly of __init__:
2 0 RESUME 0
3 2 LOAD_FAST 1 (x)
4 LOAD_FAST 0 (self)
6 STORE_ATTR 0 (x)
16 RETURN_CONST 0 (None)
This demonstrates what padding we need — and also explains why neither of our
solutions worked, and some of CACHE's underlying logic. In the first case, we
had no padding at all; in the second case, we preserved all the padding, keeping
it both after LOAD_FAST self and after STORE_ATTR x. From what we see now,
the latter should stay, but the former should not; and it stands to reason,
because some instructions can be optimized (in which case, this buffer is
needed), while others can't (in which case it's not). It seems that STORE_ATTR
can be optimized, which is why RETURN_CONST is offsetted from it by 10 zeros;
and while the same was true for LOAD_GLOBAL, it no longer applies to
LOAD_FAST, which we've replaced it with. The full reasoning is that complex
lookup operations (like globals and attributes) can be optimized by saving a
direct shortcut to the object's memory address once it's found; for simple
lookup operations, this would be immaterial, since they point to a fixed offset
already. The segmentation fault we saw earlier was actually Python's patching in
such a shortcut: since we hadn't allocated a dedicated space for it, the memory
address was written over subsequent opcodes, and when the interpreter tried to
execute code that no longer made sense, it crashed.
Unfortunately, our current solution doesn't lend itself to a quick fix that
would address this. One thing we can do is give up on our nice Instruction
objects and go back to manual parsing:
>>> remove_padding = False
>>> it = iter(code.co_code)
>>> for op, arg in zip(it, it):
... if remove_padding and op == 0:
... continue
... remove_padding = False
... opname = opcode.opname[op]
... if opname == "LOAD_GLOBAL" and code.co_names[arg] == "self":
... op, arg = opcode.opmap["LOAD_FAST"], 0
... remove_padding = True # skip subsequent padding
... elif opname == "LOAD_FAST":
... arg += 1
... new_bytecode.extend([op, arg or 0])
Wrapping something with iter and zipping it with itself is just a quick trick
to iterate over a sequence in pairs — in this case, reading two bytes at a time;
and the remove_padding flag lets us switch into a mode where we discard any
subsequent CACHEs after a LOAD_GLOBAL, while preserving them everywhere
else. This solution is already good enough; to wit:
>>> A.__init__.__code__ = code.replace(
... co_code=bytes(new_bytecode),
... co_varnames=tuple(new_varnames),
... co_nlocals=code.co_nlocals + 1,
... co_argcount=code.co_argcount + 1,
... )
>>> a = A(1)
>>> a.x
1
And it's what I'd expect someone who relies on reverse-engineering to figure
out. I add it here since I wanted to introduce you to this approach, and felt
that some closure was called for — but also, seeing as I'm writing a book and
did do some research, I thought I'd share a more appropriate solution as well.
If you'd keep digging around, you'd eventually find dis._inline_cache_entries,
which holds the number of CACHEs expected after any instruction; so we can go
back to our original implementation, and achieve the same result in a more
elegant and robust way:
>>> for i in dis.get_instructions(code):
... op, arg = i.opcode, i.arg
... if i.opname == "LOAD_GLOBAL" and i.argval == "self":
... op, arg = opcode.opmap['LOAD_FAST'], 0
... elif i.opname == "LOAD_FAST":
... arg += 1
... new_bytecode.extend([op, arg or 0])
... cache_size = dis._inline_cache_entries[op]
... if cache_size:
... new_bytecode.extend([0, 0] * cache_size)
Note that _inline_cache_entries is an internal and undocumented construct, so
it might very well change in the future. At the time of writing, however, it's
the most reasonable way to handle CACHE padding I could find. If it ever does
disappear, a more standard approach would be to use Instruction.size
(available since Python 3.11), which gives you each instruction's size – but
that version is left as an exercise to the reader.
An Implicit Self
Now that we've made it work, let's make it shine: all we need is a class
decorator that would run over all the methods and patch them to include self —
without us having had to write it. Putting our experiment in a tidy function:
>>> def add_self(method):
... code = method.__code__
... new_varnames = list(code.co_varnames)
... new_varnames.insert(0, "self")
... new_bytecode = []
... # … same loop as before
... method.__code__ = code.replace(
... co_code=bytes(new_bytecode),
... co_varnames=tuple(new_varnames),
... co_nlocals=code.co_nlocals + 1,
... co_argcount=code.co_argcount + 1,
... )
>>> def selfless(cls):
... for key, value in vars(cls).items():
... if callable(value):
... add_self(value)
... return cls
We're done! It wasn't easy — but if we're honest, it wasn't that difficult either; or at least, it all fits in rather few lines of code. The result:
>>> @selfless
... class A:
... def __init__(x):
... self.x = x
... def f():
... return self.x
>>> a = A(1)
>>> a.f()
1
Rewiring Scopes
Having done all that, it suddenly occurred to me that it was completely
unnecessary. Not the auto-injecting self bit — that was cool, I stand by it —
but the hardcore, low-level bytecode manipulation we did to achieve it.
Think about it: our problem was that Python considered self to be global,
instead of local, and we had to get our hands quite dirty to remedy this
mistake. What if instead, we'd let it continue thinking this way, and simply
patch how global variables are resolved?
This time, let's do it top-down instead of bottom-up. We want to keep the same
decorator — only now, we replace our methods with SelflessMethods, to be
defined:
>>> def selfless(cls):
... for key, value in vars(cls).items():
... if callable(value):
... setattr(cls, key, SelflessMethod(value))
... return cls
Such selfless methods, upon receiving a function, will simply give it back —
unbound, mind you, because binding it to an instance will try to squeeze an
argument into it which it does not expect. Before they do, however, they'll
replace their function's global scope, conveniently located in __globals__,
with an almost identical dictionary – albeit one where the key self is wired
to the appropriate instance.
The only complication is that since function objects are immutable, we won't be able to assign them directly, so we have to recreate them and copy over all the relevant attributes:
>>> import types
>>> class SelflessMethod:
... def __init__(self, func):
... self.func = func
... def __get__(self, instance, cls):
... return types.FunctionType(
... code=self.func.__code__,
... name=self.func.__name__,
... globals={"self": instance, **self.func.__globals__},
... argdefs=self.func.__defaults__,
... closure=self.func.__closure__,
... )
And that's basically it. Having replaced f, for example, with a
SelflessMethod, we get the following: when a.f() is called, it calls the
selfless method's __get__ to acquire a bound method, passing in the instance
a and the class A. However, instead of binding the function to the instance,
we reproduce it — same code, name, and everything else — with a slight wrinkle
in its global scope, mapping self to a. This way, when it's called, it still
thinks self is global — it just so happens that when it tries to resolve it,
it ends up with exactly the object we intended it to. Read it and weep:
>>> @selfless
... class A:
... def __init__(x):
... self.x = x
... def f():
... return self.x
>>> a = A(1)
>>> a.f()
1
So Wait, What About the Bytecode
Pedagogically speaking, this might seem like a weird maneuver: we spent a lot of time reckoning with code objects, only to discard them entirely in favor of a much simpler way. Needless to say, this was intentional – because once again, the idea is to have fun. We don't know where our adventure will take us — but if we want to wander far and wide, we shouldn't commit to one path or another. Since there wasn't a particular destination we were trying to reach, the walk we had should be its own reward — and having familiarized ourselves with new territory, we can always revisit it when it's called for. At least, that's what I meant when I was talking about "the making of people" rather than "the making of code".
Still, to give some coherent shape to our narrative, let's play with bytecode some more. To continue with the travel metaphor, I wouldn't want us to mark it on our mental map as a spooky, peripheral neighbourhood, when in actuality it's the city center, driving pretty much all of the rest of the language. True, often times we have higher-level abstractions to take care of things — but it's always there, beneath the surface (so maybe the "city sewers" would be a more accurate metaphor, but it doesn't sound as appealing).
In any case, consider the following code:
>>> def f():
... print('before')
... print('after')
>>> def g():
... print('in the middle')
If we'd want to inject some code before or after f, we could use a decorator —
but what if we want to run g in its middle? As far as I know, nothing else
in the language will let us do it, and in some cases it can be extremely useful.
For example, we might have a running system that we want to debug, like a web
server where some endpoints are misbehaving. Normally, we'd have to edit the
code to print some value, hopefully revealing the problem, or else install a
breakpoint that would let us inspect the scene — but in both cases we'd have to
restart the process. With live code injection, we can do so to any running
system without so much as a hiccup; but before we consider what such a scenario
might entail, let's solve for the simpler case of f and g.
Live Code Injection
Fortunately, this challenge is easier than the one we've already faced. There's no need to rearrange opcodes or add variables — just figure out how a function invocation looks like, and add it at the desired place. Let's reverse-engineer the first:
>>> def invocation():
... g()
>>> dis.dis(invocation)
14 0 RESUME 0
15 2 LOAD_GLOBAL 1 (NULL + g)
12 CALL 0
20 POP_TOP
22 RETURN_CONST 0 (None)
So: we bring g, which is a global name, into focus; then call it; and then pop
the top of the stack, because CALL will leave the return value there, and
since we don't use it, the least we can do is clean it up. Incidentally, if our
function would expect some arguments, you'd see all of them loaded unto the
stack before the call — but that's not something we need to worry about at the
moment.
We also see another optimization. In Python 3.10, this would be
LOAD_GLOBAL 0 (g), which makes more sense – if we check
invocation.__code__.co_names, we find g at index 0, not 1; and what is
that NULL, anyway? Turns out, it's a way to load global functions faster: it
indicates whether they expect a self (in which case we use their "correct"
index) or not (in which case we use their index plus one). Weird, but those are
the rules; just note how this implicitly "doubles" co_names, which we have to
keep in mind when indexing it ourselves.
Our main unknown is how to determine where to inject these three instructions so
they're "in the middle"; for that, we can use Instruction's starts_line:
>>> for i in dis.get_instructions(f):
... print(i.starts_line, i.opname, i.argval)
1 RESUME 0
2 LOAD_GLOBAL print
None LOAD_CONST before
None CALL 1
None POP_TOP None
3 LOAD_GLOBAL print
None LOAD_CONST after
None CALL 1
None POP_TOP None
None RETURN_CONST None
Because we're working in the interpreter, every function's first line is always
1:
>>> f.__code__.co_firstlineno
1
If we'd put it in some file, however, it would be the actual line number, and
all the starts_lines would shift accordingly. In any case, the first
instruction is always RESUME, so its starts_line is the same as
co_firstlineno, which you can roughly think of as the line on which the
function definition begins. Then, in our case, the next line contains an
invocation of print on "before", and the line after that – an invocation of
print on "after". What we want, then, is to act as soon as we see an opcode
from line 3. Well:
>>> def add_cache(opname): # remember this?
... return [0, 0] * dis._inline_cache_entries[opcode.opmap[opname]]
>>> def inject(into, line, hook):
... code = into.__code__
... new_names = list(code.co_names)
... new_names.append(hook.__name__)
... hook_index = 2 * new_names.index(hook.__name__) + 1
... new_bytecode = []
... for i in dis.get_instructions(code):
... if i.starts_line == line:
... new_bytecode.extend([
... opcode.opmap["LOAD_GLOBAL"], hook_index,
... *add_cache("LOAD_GLOBAL"),
... opcode.opmap["CALL"], 0,
... *add_cache("CALL"),
... opcode.opmap["POP_TOP"], 0,
... *add_cache("POP_TOP"),
... ])
... new_bytecode.extend([
... i.opcode,
... i.arg or 0,
... *add_cache(i.opname),
... ])
... into.__code__ = code.replace(
... co_code=bytes(new_bytecode),
... co_names=tuple(new_names),
... )
Note how this assumes that the hook function is in the target function's global scope; if this wasn't the case, we'd also need to add:
>>> def inject(into, line, hook):
... # … same as before
... into.__globals__[hook.__name__] = hook
Or, if we don't want to "pollute" its actual global namespace, which it probably shares with other objects, patch it up more sneakily, like we did with selfless methods:
>>> def inject(into, line, hook):
... # … same as before
... return types.FunctionType(
... code=into.__code__,
... name=into.__name__,
... globals={hook.__name__: hook, **into.__globals__},
... argdefs=into.__defaults__,
... closure=into.__closure__,
... )
In any case, check this out:
>>> inject(f, 3, g)
>>> f()
before
in the middle
after
Seriously now, how cool is that?