support narrowing enum values using == and !=

[DetachHead]

from typing import NoReturn, Literal
from enum import auto, Enum

def foo(value: Literal["foo","bar"]) -> None:
    if value == "foo":
        reveal_type(value) #note: Revealed type is "Literal['foo']"
    elif value != "bar":
        reveal_type(value) #error: Statement is unreachable  [unreachable]

class Foo(Enum):
    foo = auto()
    bar = auto()

def bar(value: Foo) -> None:
    if value == Foo.foo:
        reveal_type(value) #Revealed type is "__main__.Foo"
    elif value != Foo.bar:
        reveal_type(value) #no error, even though this is unreachable

https://mypy-play.net/?mypy=latest&python=3.10&flags=show-column-numbers%2Cshow-error-codes%2Cstrict%2Ccheck-untyped-defs%2Cdisallow-any-decorated%2Cdisallow-any-expr%2Cdisallow-any-explicit%2Cdisallow-any-generics%2Cdisallow-any-unimported%2Cdisallow-incomplete-defs%2Cdisallow-subclassing-any%2Cdisallow-untyped-calls%2Cdisallow-untyped-decorators%2Cdisallow-untyped-defs%2Cwarn-incomplete-stub%2Cwarn-redundant-casts%2Cwarn-return-any%2Cwarn-unreachable%2Cwarn-unused-configs%2Cwarn-unused-ignores&gist=0f58b3bb7b93c7e3d917c18a7df2ed39

[erictraut]

If you use the is and is not operators rather than == and !=, mypy will perform narrowing for enum values.

[KotlinIsland]

Why? is it an anti pattern to use ==/!= on enum values? if so should a warning be raised when they are used?.

[erictraut]

I don't know why mypy handles is and is not but not == and != in this case. I'm guessing that the former is more common than the latter. I think both patterns are reasonable. I can't think of a reason why == and != would be an anti-pattern. FWIW, pyright supports type narrowing for both pairs of operators.

[hauntsaninja]

If you want to give implementing this a shot take a look at #7000, in particular #7000 (comment)

[KotlinIsland]

@hauntsaninja Would this be as simple as changing:

mypy/mypy/checker.py

Line 4476 in 421fcc2

coerce_only_in_literal_context = True

To be False if the comparison is against an Enum?

[hauntsaninja]

Yeah, should be something like that. You'll want to restrict to when the comparison is against an Enum that has a last_known_value. Let me play around and see...

[KotlinIsland]

coerce_only_in_literal_context = False if all(isinstance(p := get_proper_type(t), Instance) and p.type.is_enum for t in expr_types) else True
I wrote this, but it looks pretty cringe.

[hauntsaninja]

This is what I'm playing with:

diff --git a/mypy/checker.py b/mypy/checker.py
index 95789831f..d8e65f73c 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -4802,7 +4802,16 @@ class TypeChecker(NodeVisitor[None], CheckerPluginInterface):
         """
         should_coerce = True
         if coerce_only_in_literal_context:
-            should_coerce = any(is_literal_type_like(operand_types[i]) for i in chain_indices)
+
+            def should_coerce_inner(typ: Type) -> bool:
+                typ = get_proper_type(typ)
+                return is_literal_type_like(typ) or (
+                    isinstance(typ, Instance)
+                    and typ.type.is_enum
+                    and (typ.last_known_value is not None)
+                )
+
+            should_coerce = any(should_coerce_inner(operand_types[i]) for i in chain_indices)
 
         target: Optional[Type] = None
         possible_target_indices = []

We do have test cases that explicitly test against this behaviour though :-( most notably testNarrowingEqualityFlipFlop.

[KotlinIsland]

Were they added when enums is comparisons were added.

To make sure that they wouldn't work on == due to complications.

[KotlinIsland]

But now that we are implementing it correctly we should change the tests.

[KotlinIsland]

and (typ.last_known_value is not None)
When would it be an enum and None?

[hauntsaninja]

It'd be the case in

def f(x: Medal, y: Medal):
    if x == y: ...

Probably fine to omit that check, since coerce_to_literal won't do anything if last_known_value is None (actually that's not true, it'll narrow singleton enums as well, so probably good to omit it, barring some chaining weirdness).

Yeah, the tests are pretty explicit about why they test what they test:

mypy/test-data/unit/check-narrowing.test

Line 720 in d41e34a

# So strictly speaking, we ought to do the same thing with 'is' comparisons

I'm not convinced that those tests are necessarily the behaviour we want, but seems like whoever wrote them collected some data / spent more time thinking about this specific issue than I have, so just saying I'd need to think about it too! :-)

[KotlinIsland]

That looks like a scenario like:

def m():
    global a
    a = 2

a: int = True

if a is True:
    m()
    reveal_type(a)  # Literal[True]

Sus!

[KotlinIsland]

What's the policy on narrowing mutable variables? seems to me like that example with narrowing to True is the same as that test?

[hauntsaninja]

Yeah, I'm inclined to agree (although literal types always make things murkier in practice). Although the test case specifically tests enums, it's possible that the real world false positives alluded to were more often narrowing to literal strs or something. I can make the PR and we'll see what other mypy maintainers have to say.

[parsons-kyle-89]

Last I checked (which was all the way back in version 0.790), the behavior was even stranger than presented here. == did narrow narrow Enums in if-else blocks, but only if they followed an is. See #9786

[KotlinIsland]

This makes match on enums really annoying.

[hauntsaninja]

I don't really use match. Could you paste a semi-realistic example where merging #11521 would help?

[KotlinIsland]

Good point, I don't know why I used a match...

But I could imagine if anyone ever does use a match that this will be annoying.