Search code examples
pythonpython-typingmypy

Type annotation for partial functions


we have multiple partial-like functions with same type annotation with args and kwargs like:

def fruit(fruit_name: str, fruit_class: Type, arg1: int, arg2: float, arg3: tuple):
    pass
    
def apple(*args, **kwargs):
    return fruit("apple", *args, **kwargs)
    
def orange(*args, **kwargs):
    return fruit("orange", *args, **kwargs)
    
def banana(*args, **kwargs):
    return fruit("banana", *args, **kwargs)

can I somehow create a template for functions apple, orange, banana and assign it to them? I thought about Protocol with __call__ definition, but it is unclear how to assign it to functions


Solution

  • Rather than a function, you can use functools.partial:

    from typing import Callable, Type
    from functools import partial
    
    Fruit = ...  # whatever fruit returns
    # type of fruit(), but without the initial string parameter
    FruitMaker = Callable[[Type, int, float, tuple], Fruit]
    
    def fruit(fruit_name: str, fruit_class: Type, arg1: int, arg2: float, arg3: tuple) -> Fruit:
        ...
    
    
    apple: FruitMaker = partial(fruit, "apple")
    orange: FruitMaker = partial(fruit, "orange")
    banana: FruitMaker = partial(fruit, "banana")
    

    Another possibility would be to refactor fruit to take a name, and return a function that closes over the fruit name.

    def fruit_maker(fruit_name: str) -> FruitMaker:
        def fruit(fruit_class: Type, arg1: int, arg2: float, arg3: tuple):
            ...
        return fruit
    
    apple = fruit_maker("apple")
    orange = fruit_maker("orange")
    banana = fruit_maker("banana")