Search code examples
pythonpython-typingmypy

How to make mypy correctly type-check a function using functools.partial?


I'm trying to create a function that returns a partially applied callable, but I'm encountering issues with mypy type checking.

HEre Is my first implementation:

Help me to explain my question for stackoverflow. i.e find a title and the body

this code :

from collections.abc import Callable
from functools import partial

def f(i: int, j: float, k: int) -> int:
    return i + int(j) + k

def g(a: float) -> Callable[[int, int], int]:
    return partial(f, j=a)

fun: Callable[[int, int], int] = g(3.0)
r: int = fun(4, 5)
print(r)

It is successfully checked by mypy

but can not run

r: int = fun(4, 5) TypeError: f() got multiple values for argument 'j'

to solve this problem, I call the function with named argument

from functools import partial

def f(i: int, j: float, k: int) -> int:
    return i + int(j) + k

def g(a: float) -> Callable[[int, int], int]:
    return partial(f, j=a)

fun: Callable[[int, int], int] = g(3.0)
# line 12 in my code (where the error message comes from)
r: int = fun(i=4, k=5)
print(r)

it works fine now

but mypy checking fails

main.py:12: error: Unexpected keyword argument "i" [call-arg] main.py:12: error: Unexpected keyword argument "k" [call-arg] Found 2 errors in 1 file (checked 1 source file)

Is there a way to annotate this code so that it both runs correctly and passes mypy's type checking? I've tried various combinations of type hints, but I haven't found a solution that satisfies both the runtime behavior and static type checking.

I know there is this solution, without using Partial

from collections.abc import Callable 


def f(i :int,j : float,k :int) ->int:
    return i+int(j)+k
    
def g(a :float) -> Callable[[int,int],int]:
    def ret(i,k):
        return f(i,a,k)
    return ret

fun :Callable[[int,int],int]= g(3.0)
r : int = fun(4,5)
print(r)

But I really want to use Callable because for I am working with functions with a lot of parameter and it this much more simplier to just say which paramters are replaced


Solution

  • Callable is not expressive enough to be able to describe the function signature you are returning. Instead, you should use a Protocol. This is caused by your use of partial() to fill in the argument for j. That is, partial(f, j=4)(1, 2) is equivalent to f(1, 2, j=4) which means Python tries to pass both 2 and 4 as the argument for j. In cannot do this, and so instead throws an error. Thus, k MUST be passed as a keyword argument instead of a positional argument.

    Instead of:

    # All this says is that the function takes two ints as argument. It does not
    # say what the name of the argument is. So you cannot use keyword arguments.
    # Even if that is what is required in the case of arg `k`
    FuncType = Callable[[int,int],int]
    

    Use:

    from typing import Protocol
    
    # Callable takes two int args called `i` and `k`. `i` can be passed as a normal
    # arg or as a keyword arg, and `k` MUST be passed as a keyword arg.
    class FuncType(Protocol):
        def __call__(self, i: int, *, k: int) -> int: ...
    

    Thus, the latter part of your code would become:

    class FuncType(Protocol):
        def __call__(self, i: int, *, k: int) -> int: ...
    
    def g(a: float) -> FuncType:
        # NB. mypy does not check this cast, even with the --strict flag
        return partial(f, j=a) 
    
    fun: FuncType = g(3.0)
    r: int = fun(4, j=5)
    # OR
    r = fun(i=4, j=5)
    # mypy is happy with either
    print(r) 
    

    Passing k as a positional argument

    It is possible to pass k as a positional argument, but not by using partial. Instead you must use a closure to wrap your calls to f. For example:

    def g(a: float) -> Callable[[int, int], int]:
        def wrapper(i: int, k: int) -> int:
            return f(i, a, k)
        return wrapper
    
    g(1.0)(2, 3) # mypy is happy with this
    

    You'll note that because k can now be passed as a positional argument you can use Callable to express the function signature again. However, even though the args of wrapper are i and k, mypy will not let you pass these args by keyword. This is because Callable does not expose the names of the arguments of the callable. If you want to be able to use keyword arguments, then you will need to use a Protocol again.