mypy icon indicating copy to clipboard operation
mypy copied to clipboard

Support enum type narrowing with `in`

Open edaqa-uncountable opened this issue 3 years ago • 0 comments

The syntax some_enum_value in [list_of_values] should narrow the enum to that list of values.

This is a list extension to https://github.com/python/mypy/issues/10915 "support narrowing enum values using == and !=" Unlike the == case, there is no equivalent is syntax for lists that will narrow the type.

This example shows the motivation for this feature, as well as the syntax that is expected to work.

from typing import Literal, Union, cast
from enum import Enum


class SomeType(Enum):
    a = 1
    b = 2
    c = 3


LimitType = Union[Literal[SomeType.a, SomeType.b]]
AllowedLimitType = [SomeType.a, SomeType.b]


def inner_check(limit: LimitType) -> None:
    pass


def outer_check(some: SomeType) -> None:
    if some in [SomeType.a, SomeType.b]:
        inner_check(some)  # Incompatible type

    if some in AllowedLimitType:
        inner_check(some)  # Incompatible type

    if some == SomeType.a or some == SomeType.b:
        inner_check(some)  # Incompatible type
        inner_check(cast(LimitType, some))

    if some is SomeType.a or some is SomeType.b:
        inner_check(some)

    if some == SomeType.a:
        inner_check(SomeType.a)

    # This would be ideal
    # if some in LimitType:
    #   inner_check(some)

def main() -> None:
    outer_check(SomeType.a)


main()

edaqa-uncountable avatar Sep 19 '22 08:09 edaqa-uncountable