from typing import Union, List, Iterable
from typing_extensions import Literal, TypedDict

# Goal here is to add mypy type so that
# if the dict's key is 'a' or 'b' mypy knows
# that the value should be int or str.
ABDictT = TypedDict('ABDictT', {
    'key': Literal['a', 'b'],
    'value': Union[str, int]
})

# If key is 'c', or 'd' then mypy should know
# value should be str, or int list.
CDDictT = TypedDict('CDDictT', {
    'key': Literal['c', 'd'],
    'value': Union[str, Iterable[int]],
})

ValueDictElementT = Union[ABDictT, CDDictT]
ValuesDictT = Iterable[ValueDictElementT]

def some_func(values: ValuesDictT) -> None:
    for element in values:
        if element['key'] == 'a':
            # The reveal_type here show the type to be
            # Union[str, int, Iterable[int]] since we
            # did a check that element['key'] is 'a' mypy
            # should be able to tell that the type is
            # ABDictT here?
            reveal_type(element['value'])