How to wrap a function with accurate __code__.co_argcount?

After a lot of tinkering, I found a pretty procedural way that might work for you.

The trick was to use __code__.replace(). There are some caveats, probably more than I know.

def high_order_function(f):
    """high order function expecting a function with 3 arguments!"""

    print(f"\nprocessing function {f.__name__}")
    if f.__code__.co_argcount!=3:
        raise Exception(f"Error Expecting a function with 3 arguments, the passed function got {f.__code__.co_argcount}")
    print("Function is Ok")

    #...
    return None

def my_wrapper(func):
    import functools

    @functools.wraps(func)
    def inner(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w, x, y, z):
        kwargs = locals().copy()
        del kwargs["func"]
        print("wrapped1!")
        r = func(**kwargs)  # func(*kwargs.values()) would work too
        print("wrapped2!")
        return r

    func_args = func.__code__.co_varnames
    inner.__code__ = inner.__code__.replace(co_varnames=func_args, co_argcount=len(func_args))

    return inner

def original(a, b, c):
    return None

wrapped = my_wrapper(original)

high_order_function(original)
high_order_function(wrapped)

Result

processing function original
Function is Ok

processing function original
Function is Ok

functools.wraps changes the name of inner to original

Caveats

• __code__.replace() raised ValueError: code: varnames is too small when inner's parameters were *args or **kwargs
• If inner instead had no parameters then locals() inside it would not get the supplied values, therefore you got the whole alphabet instead
• Inside inner you can access the parameters by the letter like normal if you're sure it´s supplied otherwise you'll get IndexError: tuple index out of range
• I recommend to use e.g. kwargs.get("d") instead
• __code__.replace may only be for 3.8+, it has sys.version_info >= (3, 8) in the source code

Why not define 2 version of the inner function in your wrapper based on whether you need to pass the wrapped function to higher_order_function or not.

Something like this:

def high_order_function(f):
    """high order function expecting a function with 3 arguments!"""

    print(f"\nprocessing function {f.__name__}")
    if f.__code__.co_argcount != 3:
        raise Exception(
            f"Error Expecting a function with 3 arguments, the passed function got {f.__code__.co_argcount}")
    print("Function is Ok")

    # ...
    return None


def my_wrapper(func, higher_order_compatible=True):  # switch to control whether should be compatible with the higher_order_function or not
    import functools

    if higher_order_compatible:
        @functools.wraps(func)
        def inner(a, b, c, *args, **kwargs):  # extra args to satisfy the condition
            print("wrapped1!")
            r = func(a, b, c, *args, **kwargs)
            print("wrapped2!")
            return r
    else:
        @functools.wraps(func)
        def inner(*args, **kwargs):  # normal wrapper
            print("wrapped1!")
            r = func(*args, **kwargs)
            print("wrapped2!")
            return r

    return inner


def original(a, b, c):
    return None


wrapped_compat = my_wrapper(original, higher_order_compatible=True)
wrapped_nocompat = my_wrapper(original, higher_order_compatible=False)

print("Original")
high_order_function(original)
print("Compatible")
high_order_function(wrapped_compat)
print("Not compatible")
try:
    high_order_function(wrapped_nocompat)
except:
    print("no not working")

Results in:

Original
processing function original
Function is Ok

Compatible
processing function original
Function is Ok

Not compatible
processing function original
no not working

co_argcount: number of arguments (not including keyword only arguments, * or ** args)

Hence the goal is to bypass such definition. Make a fake signature with 3 fake parameters, these are taken into consideration by the code attribute co_argcount. Then the parameters of the original function must by passed as keys.

def wrapper(f):
    def extended_signature(fake1=None, fake2=None, fake3=None, **kwargs):
        return f(**kwargs)
    return extended_signature


def a(q, w): print(q, w)

a_wrapped = wrapper(a)

high_order_function(a_wrapped)(q=1, w=2)
#processing function true_signature
#Function is Ok

This is an attempt to solve the problem

i feel like it's almost a potential solution

However it is not working as expected, strange. f()for a function object class is not as f.__call__() hmm

def my_wrapper(func):
    import copy

    def inner(*args, **kwargs):
        print("wrapped front")
        r = func(*args, **kwargs)
        print("wrapped end")
        return r

    newfunc = copy.deepcopy(func)
    newfunc.__name__ = func.__name__ + "_wrapped"
    newfunc.__call__ = inner

    return newfunc

def original(a, b, c=6):
    print("original",a,b,c)
    return None

###testing if the original function work
high_order_function(original)
#will pass requirement

###testing if the wrap works?
high_order_function(my_wrapper(original))
#will pass requirement, however the wrap did not work

well, here's a solution, couldn't find a procedural way to generate the functions... it scales up to 5 forced arguments

def my_wrapper(func):
    """see https://stackoverflow.com/questions/73601340/how-to-wrap-a-function-with-accurate-code-argcount?noredirect=1#comment129973896_73601340
    yes this is a shit show, did not found a procedural way to generate functions. tried exec() code generation & was also a mess"""

    #find back the expected arguments so func.__code__.co_argcount will be accurate
    if (func.__defaults__ is not None):
          force_arg = func.__code__.co_argcount - len(func.__defaults__)
    else: force_arg = func.__code__.co_argcount

    import functools

    if (force_arg==0):
            
        @functools.wraps(func)
        def inner(**kwargs):
            print("wrapped1!")
            r = func(**kwargs)
            print("wrapped2!")
            return r

    elif (force_arg==1):

        @functools.wraps(func)
        def inner(a,**kwargs):
            print("wrapped1!")
            r = func(a,**kwargs)
            print("wrapped2!")
            return r

    elif (force_arg==2):

        @functools.wraps(func)
        def inner(a,b,**kwargs):
            print("wrapped1!")
            r = func(a,b,**kwargs)
            print("wrapped2!")
            return r

    elif (force_arg==3):

        @functools.wraps(func)
        def inner(a,b,c,**kwargs):
            print("wrapped1!")
            r = func(a,b,c,**kwargs)
            print("wrapped2!")
            return r

    elif (force_arg==4):

        @functools.wraps(func)
        def inner(a,b,c,d,**kwargs):
            print("wrapped1!")
            r = func(a,b,c,d,**kwargs)
            print("wrapped2!")
            return r

    elif (force_arg==5):

        @functools.wraps(func)
        def inner(a,b,c,d,e,**kwargs):
            print("wrapped1!")
            r = func(a,b,c,d,e,**kwargs)
            print("wrapped2!")
            return r

    else: raise Exception("my_wrapper() do not support more than 5 forced argument")

    return inner