How to type-hint a function return based on input parameter value?

user_5811400 0
Viewed 334

How can I type-hint a function in Python based on the value of an input parameter?

For instance, consider the following snippet:

from typing import Iterable

def build(
    source: Iterable,
    factory: type
) -> ?: # what can I write here?
    return factory(source)

as_list = build('hello', list) # -> list ['h', 'e', 'l', 'l', 'o']
as_set = build('hello', set) # -> set {'h', 'e', 'l', 'o'}

When building as_list, the value of factory is list, and this should be the type annotation.

I am aware of this other question, but, in that case, the return type depended only on the input types, not on their values. I would like to have def build(source: Iterable, factory: type) -> factory, but of course this doesn't work.

I am also aware of Literal types in Python 3.8+, and something similar to this could be achieved:

from typing import Iterable, Literal, overload
from enum import Enum

FactoryEnum = Enum('FactoryEnum', 'LIST SET')

@overload
def build(source: Iterable, factory: Literal[FactoryEnum.LIST]) -> list: ...

@overload
def build(source: Iterable, factory: Literal[FactoryEnum.SET]) -> set: ...

But this solution would make factory useless (I could just define two functions build_list(source) -> list and build_set(source) -> set).

How can this be done?

type-hinting python-3.x python
1 Answers

Rather than using type, you could use a generic and define the factory as a Callable, as follows:

from typing import Callable, Iterable, TypeVar

T = TypeVar('T')

def build(
    source: Iterable,
    factory: Callable[[Iterable], T]
) -> T:
    return factory(source)
user_30017610
Related