CODEWITHSUNDEEP
pythonmypypython-typing
Vringar
Vringar

Reputation: 572

How to assure mypy that user input is in a set of literals

My code is basically the following:

def foo(param: Literal["a", "b"]) -> None:
    ...

i = input("A or B?").lower()
assert i in ["a","b"]
foo(i)

However mypy complains that it can't know that

Argument 1 to "foo" has incompatible type "str"; expected "Union[Literal['a'], Literal['b']]"

and I can't find documentation on how to convert this to a literal expression.

Upvotes: 2

Views: 346

Answers (1)

Yaakov Saxon
Yaakov Saxon

Reputation: 378

As @juanpa.arrivillaga mentioned in a comment, you can use typing.cast to cast the value as an element. Unfortunately mypy doesn't appear to do any checking on type.cast that the value actually fits the literal it is being cast to, but you can do that check for that yourself against typing.get_args.

a_or_b=Literal["a", "b"]
def foo(param: a_or_b) -> None:
    ...


i = input("A or B?").lower()
if not i in typing.get_args(a_or_b):
    raise ValueError()
else:
    i = typing.cast(i,a_or_b)
foo(i)

Upvotes: 2

Related Questions