Starting point

def old(foo, bar):
    """This is old's docstring."""
    print(foo, bar)
    return foo + bar


def new(prefix, foo, *args, **kwargs):
    return old(prefix + foo, *args, **kwargs)

Let’s test it.

>>> o = old('a', 'b')
a b
>>> n = new('!', 'a', 'b')
!a b
>>> print(o, n, sep=' - ')
ab - !ab
>>> help(old)
Help on function old in module __main__:

old(foo, bar)
    This is old's docstring.

>>> help(new)
Help on function new in module __main__:

new(prefix, foo, *args, **kwargs)

The last line is not exactly informative — it doesn’t tell us that we need to pass bar as an argument. Sure, you could define new as just (prefix, foo, bar) — but that means every change to old requires editing new as well. So, not ideal. Let’s try to fix this.

The existing infrastructure: functools.wraps

First, let’s start with the basic facility Python already has. The standard library already comes with functools.wraps and functools.update_wrapper.

If you’ve never heard of those two functions, here’s a crash course:

def decorator(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        print("Inside wrapper")
        f(*args, **kwargs)
    return wrapper

@decorator
def square(n: float) -> float:
    """Square a number."""
    return n * n

If we try to inspect the square function, we’ll see the original name, arguments, annotations, and the docstring. If we ran this code again, but with the @functools.wraps(f) line commented out, we would only see wrapper(*args, **kwargs).

This approach gives us a hint of what we need to do. However, if we apply wraps (or update_wrapper, which is what wraps ends up calling) to our function, it will only have foo and bar as arguments, and its name will be displayed as old.

So, let’s take a look at functools.update_wrapper. What does it do? Two things:

• copy some attributes from the old function to the new one (__module__, __name__, __qualname__, __doc__, __annotations__)

• update __dict__ of the new function

• set wrapper.__wrapped__

If we try to experiment with it — by changing the list of things to copy, for example — we’ll find out that the annotations, the docstring, and the displayed name come from the copied attributes, but the signature itself is apparently taken from __wrapped__.

Further investigation reveals this fact about inspect.signature:

inspect.signature(callable, *, follow_wrapped=True)

New in version 3.5: follow_wrapped parameter. Pass False to get a signature of callable specifically (callable.__wrapped__ will not be used to unwrap decorated callables.)

And so, this is our end goal:

Craft a function with a specific signature (that merges old and new) and set it as new.__wrapped__.

But first, we need to talk about parallel universes.

Or actually, code objects.

Defining a function programmatically

Let’s try an experiment.

>>> def foo(bar): pass
>>> foo.__wrapped__ = lambda x, y: None
>>> help(foo)
foo(x, y)

So, there are two ways to do this. The first one would be to generate a string with the signature and just use eval to get a __wrapped__ function. But that would be cheating, and honestly, quite boring. (The inspect module could help us with preparing the string.) The second one? Create code objects manually.

code(argcount, kwonlyargcount, nlocals, stacksize, flags, codestring,
    constants, names, varnames, filename, name, firstlineno,
    lnotab[, freevars[, cellvars]])

Create a code object.  Not for the faint of heart.

>>> def f(a, b=1, c=2, *, d=3): pass
>>> inspect.getfullargspec(f)
FullArgSpec(args=['a', 'b', 'c'], varargs=None, varkw=None, defaults=(1, 2), kwonlyargs=['d'], kwonlydefaults={'d': 3}, annotations={})

Final results

Before I show you the code, let’s test it out:

# old defined as before
@merge_args(old)
def new(prefix, foo, *args, **kwargs):
    return old(prefix + foo, *args, **kwargs)

And the end result — help(new) says:

new(prefix, foo, bar)
    This is old's docstring.

We did it!

The code is available on GitHub and on PyPI (pip install merge_args). There’s also an extensive test suite.

PS. you might be interested in another related post of mine, in which I reverse-engineer the compilation of a function: Gynvael’s Mission 11 (en): Python bytecode reverse-engineering

Reverse-engineering Python bytecode: re-creating the function by hand

I started by recreating the original firmware file. I created an empty function and wrote some code to print out __code__ contents and dis.dis output. I also added color-coding to help me read it:

#!/usr/bin/env python2
import dis
import sys

# Write code here
def check_password(s):
    pass

# Reverse engineering the code
cnames = ('co_argcount', 'co_consts', 'co_flags', 'co_name', 'co_names', 'co_nlocals', 'co_stacksize', 'co_varnames')
cvalues = (1, (None, '4e5d4e92865a4e495a86494b5a5d49525261865f5758534d4a89', 'hex', 89, 255, 115, 50), 67, 'check_password', ('decode', 'len', 'False', 'all', 'zip', 'ord'), 4, 6, ('s', 'good', 'cs', 'cg'))

for n, ov in zip(cnames, cvalues):
    v = getattr(check_password.__code__, n)
    if v == ov:
        sys.stderr.write('\033[1;32m')
    else:
        sys.stderr.write('\033[1;31m')
    sys.stderr.flush()

    sys.stdout.write(str(n) + " " + str(v) + "\n")
    sys.stdout.flush()

    sys.stderr.write('\033[0m')
    sys.stderr.flush()

dis.dis(check_password)

If we run this solver, we get the following output (text in brackets added by me):

co_argcount 1            [OK]
co_consts (None,)        [1/7 match]
co_flags 67              [OK]
co_name check_password   [OK]
co_names ()              [0/6 match]
co_nlocals 1             [should be 4]
co_stacksize 1           [should be 6]
co_varnames ('s',)       [1/4 match]
  7           0 LOAD_CONST               0 (None)
              3 RETURN_VALUE

We can see (with the help of colors, not reproduced here), that we’ve got co_argcount, co_flags, co_name correctly. We also have one constant (None, in every function) and one variable name (s, the argument name). We can also see dis.dis() output. While it looks similar to the assignment, there are a few noticeable differences: there is no 7 (line number) at the start, and LOAD_CONST instructions in the original code did not have anything in parentheses (only comparisions and loops did). This makes reading bytecode harder, but still possible. (I originally thought about using diff for help, but it’s not hard to do it by hand. I did use diff for the final checking after a manual “conversion”)

Let’s stop to look at the constants and names for a second. The long string is followed by hex, and one of the constants is decode. This means that we need to use str.decode('hex') to create a (byte)string of some information. Puzzle answers tend to be human-readable, and this string isn’t — so we need to do some more work.

So, let’s try reproducing the start of the original mission code using what we’ve just discussed. Python’s VM is based on a stack. In the bytecode above, you can see that instructions take 0 or 1 arguments. Some of them put things on the stack, others do actions and remove them. Most instruction names are self-explanatory, but the full list can be found in the dis module documentation.

Instructions like LOAD and STORE refer to indices in the constants/names/varnames tuples. To make it easier, here’s a “table” of them:

constants
 0     1                                                       2      3   4    5    6
(None, '4e5d4e92865a4e495a86494b5a5d49525261865f5758534d4a89', 'hex', 89, 255, 115, 50)

names (globals, attributes)
 0         1      2        3      4      5
('decode', 'len', 'False', 'all', 'zip', 'ord')

varnames (locals, _fast)
 0    1       2     3
('s', 'good', 'cs', 'cg')

In order to improve readability, I will use “new” dis output with names in parentheses below:

 0 LOAD_CONST               1 ('4e5d4e92865a4e495a86494b5a5d49525261865f5758534d4a89')
 3 LOAD_ATTR                0 (decode)
 6 LOAD_CONST               2 ('hex')
 9 CALL_FUNCTION            1 # function takes 1 argument from stack
12 STORE_FAST               1 (good)

As I guessed before, the first line of our function is as follows:

def check_password(s):
    good = '4e5d4e92865a4e495a86494b5a5d49525261865f5758534d4a89'.decode('hex')  # new

If we run the solver again, we’ll see that the first 12 bytes of our bytecode match the mission text. We can also see that varnames is filled in half, we’ve added two constants, and one name. The next few lines are as follows:

15 LOAD_GLOBAL              1
18 LOAD_FAST                0
21 CALL_FUNCTION            1
24 LOAD_GLOBAL              1
27 LOAD_FAST                1
30 CALL_FUNCTION            1
33 COMPARE_OP               3 (!=)
36 POP_JUMP_IF_FALSE       43
39 LOAD_GLOBAL              2
42 RETURN_VALUE

We can see that we’re putting a global object on stack and calling it with one argument. In both cases, the global has the index 1, that’s len. The two arguments are s and good. We put both lengths on stack, then compare them. If the comparison fails (they’re equal), we jump to the instruction starting at byte 43, otherwise we continue execution to load the second global (False) and return it. This wall of text translates to the following simple code:

def check_password(s):
    good = '4e5d4e92865a4e495a86494b5a5d49525261865f5758534d4a89'.decode('hex')
    if len(s) != len(good):  # new
        return False         # new

Let’s take another look at our names. We can see we’re missing all, zip, ord. You can already see a common pattern here: we will iterate over both strings at once (using zip), do some math based on the character’s codes (ord), and then check if all results (of a comparison, usually) are truthy.

Here’s the bytecode with value annotations and comments, which explain what happens where:

>>   43 LOAD_GLOBAL              3 (all)
     46 BUILD_LIST               0
     49 LOAD_GLOBAL              4 (zip)
     52 LOAD_FAST                0 (s)
     55 LOAD_FAST                1 (good)
     58 CALL_FUNCTION            2           # zip(s, good)
     61 GET_ITER                             # Start iterating: iter()
>>   62 FOR_ITER                52 (to 117)  # for loop iteration start (if iterator exhausted, jump +52 bytes to position 117)
     65 UNPACK_SEQUENCE          2           # unpack a sequence (a, b = sequence)
     68 STORE_FAST               2 (cs)      # cs = item from s
     71 STORE_FAST               3 (cg)      # cg = item from good
     74 LOAD_GLOBAL              5 (ord)
     77 LOAD_FAST                2 (cs)
     80 CALL_FUNCTION            1           # put ord(cs) on stack
     83 LOAD_CONST               3 (89)
     86 BINARY_SUBTRACT                      # - 89   [subtract 89 from topmost value]
     87 LOAD_CONST               4 (255)
     90 BINARY_AND                           # & 255  [bitwise AND with topmost value]
     91 LOAD_CONST               5 (115)
     94 BINARY_XOR                           # ^ 115  [bitwise XOR with topmost value]
     95 LOAD_CONST               6 (50)
     98 BINARY_XOR                           # ^ 50   [bitwise XOR with topmost value]
     99 LOAD_GLOBAL              5 (ord)
    102 LOAD_FAST                3 (cg)
    105 CALL_FUNCTION            1           # put ord(cs) on stack
    108 COMPARE_OP               2 (==)      # compare the two values on stack
    111 LIST_APPEND              2           # append topmost value to the list in topmost-1; pop topmost (append to list created in comprehension)
    114 JUMP_ABSOLUTE           62           # jump back to start of loop
>>  117 CALL_FUNCTION            1           # after loop: call all([list comprehension result])
    120 RETURN_VALUE                         # return value returned by all()

We can now write the full answer.

listings/gynvaels-mission-11-en/mission11.py (Source)

In the end, our dis.dis() output matches the mission text (except the removed values, but their IDs do match), our co_* variables are all green, and we can get to work on solving the puzzle itself!

Side note: this task uses a list comprehension. You might want to optimize it, remove the brackets, and end up with a generator expression. This would make the task harder, since would require working with the internal generator code object as well:

co_consts (None, '4e5d4e92865a4e495a86494b5a5d49525261865f5758534d4a89', 'hex', <code object <genexpr> at 0x104a86c30, file "mission11-genexpr.py", line 11>)

46 LOAD_CONST               3 (<code object <genexpr> at 0x104a86c30, file "mission11-genexpr.py", line 11>)

BINARY_* and ord disappeared from the new listing. You can see the modified code (which differs by two bytes) and solver output.

Solving the real puzzle

I solved the extra credit part of the puzzle. The real aim of the puzzle was to recover the password — the text for which check_password() will return True.

This part is pretty boring. I built a dictionary, where I mapped every byte (0…255) to the result of the calculation done in the check_password() function’s loop. Then I used that to recover the original text.

pass_values = {}
for i in range(256):
    result = i - 89 & 255 ^ 115 ^ 50
    pass_values[result] = i

good = '4e5d4e92865a4e495a86494b5a5d49525261865f5758534d4a89'.decode('hex')
password = ''
for c in good:
    password += chr(pass_values[ord(c)])

print(password)
print(check_password(password))

The password is: huh, that actually worked!.

What was that Paint thing about?

Yesterday’s mission was about Elvish — I mean Paint — I mean Python programming.

Most of my readers were probably puzzled by the mention of Paint. Long-time viewers of Gynvael’s streams in Polish remember the Python 101 video he posted on April Fools last year. See original video, explanation, code (video and explanation are both Polish; you can get the gist of the video without hearing the audio commentary though.) Spoilers ahead.

In that prank, Gynvael taught Python basics. The first part concerned itself with writing bytecode by hand. The second part (starts around 12:00) was about drawing custom Python modules. In Paint. Yes, Paint, the simple graphics program included with Microsoft Windows. He drew a custom Python module in Paint, and saved it using the BMP format. It looked like this (zoomed PNG below; download gynmod.bmp):

How was this done? There are three things that come into play:

• Python can import modules from a ZIP file (if it’s appended to sys.path). Some tools that produce .exe files of Python code use this technique; the old .egg file format also used ZIPs this way.

• BMP files have their header at the start of a file.

• ZIP files have their header at the end of a file.

• Thus, one file can be a valid BMP and ZIP at the same time.

I took the code of check_password and put it in mission11.py (which I already cited above). Then I compiled to .pyc and created a .zip out of it.

listings/gynvaels-mission-11-en/mission11.py (Source)

Since I’m not an expert in any of the formats, I booted my Windows virtual machine and blindly copied the parameters used by Gynvael to open the ZIP file (renamed .raw) in IrfanView and saved as .bmp. I changed the size to 83×2, because my ZIP file was 498 bytes long (3 BPP * 83 px * 2 px = 498 bytes) — by doing that, and through sheer luck with the size, I could avoid adding comments and editing the ZIP archive. I ended up with this (PNG again; download mission11.bmp):

The .bmp file is runnable! We can use this code:

listings/gynvaels-mission-11-en/ziprunner.py (Source)

And we get this:

Resources

• mission11-solver.py (full solver code)

• mission11-genexpr.py, mission11-genexpr.txt (used for side note regarding generator expressions vs. list comprehensions)

• mission11.py code, used in BMP file

• ziprunner.py, file that runs the BMP/ZIP module (adapted from Gynvael’s)

• gynmod.bmp

• mission11.bmp

• dis module documentation.

Thanks for the mission (and BMP idea), Gynvael!