operand_types = []

for expr in node.operands:

if expr not in type_map:

return {}, {}

operand_types.append(coerce_to_literal(type_map[expr]))

type_maps = []

for i, (operator, left_expr, right_expr) in enumerate(node.pairwise()):

left_type = operand_types[i]

right_type = operand_types[i + 1]

if_map = {} # type: TypeMap

else_map = {} # type: TypeMap

if operator in {'in', 'not in'}:

right_item_type = builtin_item_type(right_type)

if right_item_type is None or is_optional(right_item_type):

continue

if (isinstance(right_item_type, Instance)

and right_item_type.type.fullname() == 'builtins.object'):

continue

if (is_optional(left_type) and literal(left_expr) == LITERAL_TYPE

and not is_literal_none(left_expr) and

is_overlapping_erased_types(left_type, right_item_type)):

if_map, else_map = {left_expr: remove_optional(left_type)}, {}

else:

continue

elif operator in {'==', '!='}:

if_map, else_map = self.narrow_given_equality(

left_expr, left_type, right_expr, right_type, assume_identity=False)

elif operator in {'is', 'is not'}:

if_map, else_map = self.narrow_given_equality(

left_expr, left_type, right_expr, right_type, assume_identity=True)

else:

continue

if operator in {'not in', '!=', 'is not'}:

if_map, else_map = else_map, if_map

type_maps.append((if_map, else_map))

if len(type_maps) == 0:

return {}, {}

elif len(type_maps) == 1:

return type_maps[0]

else:

# Comparisons like 'a == b == c is d' is the same thing as

# '(a == b) and (b == c) and (c is d)'. So after generating each

# individual comparison's typemaps, we "and" them together here.

# (Also see comments below where we handle the 'and' OpExpr.)

final_if_map, final_else_map = type_maps[0]

for if_map, else_map in type_maps[1:]:

final_if_map = and_conditional_maps(final_if_map, if_map)

final_else_map = or_conditional_maps(final_else_map, else_map)

return final_if_map, final_else_map

def narrow_given_equality(self,

left_expr: Expression,

left_type: Type,

right_expr: Expression,

right_type: Type,

assume_identity: bool,

) -> Tuple[TypeMap, TypeMap]:

"""Assuming that the given 'left' and 'right' exprs are equal to each other, try

# For the sake of simplicity, we currently attempt inferring a more precise type

# for just one of the two variables.

comparisons = [

(left_expr, left_type, right_type),

(right_expr, right_type, left_type),

]

for expr, expr_type, other_type in comparisons:

# The 'expr' isn't an expression that we can refine the type of. Skip

# attempting to refine this expr.

if literal(expr) != LITERAL_TYPE:

continue

# Case 1: If the 'other_type' is a singleton (only one value has

# the specified type), attempt to narrow 'expr_type' to just that

# singleton type.

if is_singleton_type(other_type):

if isinstance(other_type, LiteralType) and other_type.is_enum_literal():

if not assume_identity:

# Our checks need to be more conservative if the operand is

# '==' or '!=': all bets are off if either of the two operands

# has a custom `__eq__` or `__ne__` method.

#

# So, we permit this check to succeed only if 'other_type' does

# not define custom equality logic

if not uses_default_equality_checks(expr_type):

continue

if not uses_default_equality_checks(other_type.fallback):

continue

fallback_name = other_type.fallback.type.fullname()

expr_type = try_expanding_enum_to_union(expr_type, fallback_name)

target_type = [TypeRange(other_type, is_upper_bound=False)]

return conditional_type_map(expr, expr_type, target_type)

# Case 2: Given expr_type=Union[A, None] and other_type=A, narrow to just 'A'.

#

# Note: This check is actually strictly speaking unsafe: stripping away the 'None'

# would be unsound in the case where A defines an '__eq__' method that always

# returns 'True', for example.

#

# We implement this check partly for backwards-compatibility reasons and partly

# because those kinds of degenerate '__eq__' implementations are probably rare

# enough that this is fine in practice.

#

# We could also probably generalize this block to strip away *any* singleton type,

# if we were fine with a bit more unsoundness.

if is_optional(expr_type) and not is_optional(other_type):

if is_overlapping_erased_types(expr_type, other_type):

return {expr: remove_optional(expr_type)}, {}

return {}, {}

def uses_default_equality_checks(typ: Type) -> bool:

"""Returns 'true' if we know for certain that the given type is using

if isinstance(typ, UnionType):

return all(map(uses_default_equality_checks, typ.items))

# TODO: Generalize this so it'll handle other types with fallbacks

if isinstance(typ, LiteralType):

typ = typ.fallback

if isinstance(typ, Instance):

typeinfo = typ.type

eq_sym = typeinfo.get('__eq__')

ne_sym = typeinfo.get('__ne__')

if eq_sym is None or ne_sym is None:

return False

return (eq_sym.fullname == 'builtins.object.__eq__'

and ne_sym.fullname == 'builtins.object.__ne__')

else:

return False

assert len(operators) + 1 == len(operands)

def pairwise(self) -> Iterator[Tuple[str, Expression, Expression]]:

"""If this comparison expr is "a < b is c == d", yields the sequence

for i, operator in enumerate(self.operators):

yield operator, self.operands[i], self.operands[i + 1]

[case testEnumReachabilityChecksBasicIdentity]

[case testEnumReachabilityChecksBasicEquality]

from enum import Enum

from typing_extensions import Literal

class Foo(Enum):

A = 1

B = 2

C = 3

x: Literal[Foo.A, Foo.B, Foo.C]

if x == Foo.A:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.A]'

elif x == Foo.B:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.B]'

elif x == Foo.C:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.C]'

else:

reveal_type(x) # No output here: this branch is unreachable

if Foo.A == x:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.A]'

elif Foo.B == x:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.B]'

elif Foo.C == x:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.C]'

else:

reveal_type(x) # No output here: this branch is unreachable

y: Foo

if y == Foo.A:

reveal_type(y) # N: Revealed type is 'Literal[__main__.Foo.A]'

elif y == Foo.B:

reveal_type(y) # N: Revealed type is 'Literal[__main__.Foo.B]'

elif y == Foo.C:

reveal_type(y) # N: Revealed type is 'Literal[__main__.Foo.C]'

else:

reveal_type(y) # No output here: this branch is unreachable

if Foo.A == y:

reveal_type(y) # N: Revealed type is 'Literal[__main__.Foo.A]'

elif Foo.B == y:

reveal_type(y) # N: Revealed type is 'Literal[__main__.Foo.B]'

elif Foo.C == y:

reveal_type(y) # N: Revealed type is 'Literal[__main__.Foo.C]'

else:

reveal_type(y) # No output here: this branch is unreachable

[builtins fixtures/bool.pyi]

[case testEnumReachabilityDisabledGivenCustomEquality]

from typing import Union

from enum import Enum

class Parent(Enum):

def __ne__(self, other: object) -> bool: return True

class Foo(Enum):

A = 1

B = 2

def __eq__(self, other: object) -> bool: return True

class Bar(Parent):

A = 1

B = 2

class Ok(Enum):

A = 1

B = 2

x: Foo

if x is Foo.A:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.A]'

else:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.B]'

if x == Foo.A:

reveal_type(x) # N: Revealed type is '__main__.Foo'

else:

reveal_type(x) # N: Revealed type is '__main__.Foo'

y: Bar

if y is Bar.A:

reveal_type(y) # N: Revealed type is 'Literal[__main__.Bar.A]'

else:

reveal_type(y) # N: Revealed type is 'Literal[__main__.Bar.B]'

if y == Bar.A:

reveal_type(y) # N: Revealed type is '__main__.Bar'

else:

reveal_type(y) # N: Revealed type is '__main__.Bar'

z1: Union[Bar, Ok]

if z1 is Ok.A:

reveal_type(z1) # N: Revealed type is 'Literal[__main__.Ok.A]'

else:

reveal_type(z1) # N: Revealed type is 'Union[__main__.Bar, Literal[__main__.Ok.B]]'

z2: Union[Bar, Ok]

if z2 == Ok.A:

reveal_type(z2) # N: Revealed type is 'Union[__main__.Bar, __main__.Ok]'

else:

reveal_type(z2) # N: Revealed type is 'Union[__main__.Bar, __main__.Ok]'

[builtins fixtures/primitives.pyi]

[case testEnumReachabilityWithChaining]

from enum import Enum

class Foo(Enum):

A = 1

B = 2

x: Foo

y: Foo

if x is Foo.A is y:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.A]'

reveal_type(y) # N: Revealed type is 'Literal[__main__.Foo.A]'

else:

reveal_type(x) # N: Revealed type is '__main__.Foo'

reveal_type(y) # N: Revealed type is '__main__.Foo'

if x == Foo.A == y:

reveal_type(x) # N: Revealed type is 'Literal[__main__.Foo.A]'

reveal_type(y) # N: Revealed type is 'Literal[__main__.Foo.A]'

else:

reveal_type(x) # N: Revealed type is '__main__.Foo'

reveal_type(y) # N: Revealed type is '__main__.Foo'

[builtins fixtures/primitives.pyi]