How to create dynamically decorated (wrapped) functions that are pickable in python?

Question

Basically, how to make the second one work. The use case is to wrap function and trap exceptions/add timings etc.

import concurrent.futures
import functools
def with_print(func):
    """ Decorate a function to print its arguments.
    """
    @functools.wraps(func)
    def my_func(*args, **kwargs):
        print("LOOK", args, kwargs)
        return func(*args, **kwargs)
    return my_func

def f():
    print('f called')

g = with_print(f)

executor = concurrent.futures.ProcessPoolExecutor(max_workers=10)

tasks = [f for x in range(10)]
fut = list()
for task in tasks:
    fut.append(executor.submit(task))
res = [x.result() for x in fut]
print(res)


# THIS ONE FAILS
tasks = [g for x in range(10)]
fut = list()
for task in tasks:
    fut.append(executor.submit(task))
res = [x.result() for x in fut]
print(res)

Error is:

_pickle.PicklingError: Can't pickle : it's not the same object as main.f

Answer 1

Define the inner function outside the decorator function and use the fact that functools.partial is pickleable:

import concurrent.futures
import functools


def inner_with_print(*args, func=None, **kwargs):
    print("LOOK", args, kwargs)
    return func(*args, **kwargs)


def with_print(func):
    result_func = functools.partial(inner_with_print, func=func)
    return functools.wraps(func)(result_func)


def f(arg, kwarg):
    print("f called")


g = with_print(f)

if __name__ == "__main__":
    with concurrent.futures.ProcessPoolExecutor(max_workers=10) as executor:
        [executor.submit(g, i, kwarg=i) for i in range(10)]

# LOOK (0,) {'kwarg': 0}
# called
# LOOK (1,) {'kwarg': 1}
# f called
# ...

Metadata was copied correctly:

print(vars(g))
# {'__module__': '__main__', '__name__': 'f', '__qualname__': 'f', '__doc__': None, '__annotations__': {}, '__wrapped__': <function f at 0x7f3251f01430>}

EDIT: the above works, but it looks like it's not an issue with dynamic decorators. Everything works fine if you change g = with_print(f) to f = with_print(f). It looks like pickle looks for __main__.f dynamically and it finds g, as a result of functools.wraps magic.

EDIT2: the functools.wraps magic is actually setting __qualname__ to f. If you set it back to g then it works fine:

g.__qualname__ = "g"

It looks like all of this is happening, because you used wraps, but also changed the function name to g.