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

Question

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.

Answer 1

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)