Skip to content

Commit

Permalink
Make enum values final - fixes #11919 (#11962)
Browse files Browse the repository at this point in the history
This fix allows using enum values in place of literals.

For example:

```
def is_a(a: Literal["a"]): return None
class MyStr(Enum):
    a = "a"
    b = "b"

is_a(MyStr.a.value)
```

Currently this fails with

```
error: Argument 1 to "is_a" has incompatible type "str"; expected "Literal['a']"
```

The fix itself is simple - the special casing of  final status for enums was being
called too late to have an effect. Moving the function call up solves
the problem.

Fixes #11919
  • Loading branch information
joey-laminar authored Jan 10, 2022
1 parent 9ea6d92 commit 48d810d
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 19 deletions.
2 changes: 1 addition & 1 deletion mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -2104,10 +2104,10 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
s.is_final_def = self.unwrap_final(s)
self.analyze_lvalues(s)
self.check_final_implicit_def(s)
self.store_final_status(s)
self.check_classvar(s)
self.process_type_annotation(s)
self.apply_dynamic_class_hook(s)
self.store_final_status(s)
if not s.type:
self.process_module_assignment(s.lvalues, s.rvalue, s)
self.process__all__(s)
Expand Down
45 changes: 30 additions & 15 deletions test-data/unit/check-enum.test
Original file line number Diff line number Diff line change
Expand Up @@ -618,8 +618,8 @@ reveal_type(A2.x.value) # N: Revealed type is "builtins.int"
reveal_type(A2.x._value_) # N: Revealed type is "builtins.int"
is_x(reveal_type(A3.x.name)) # N: Revealed type is "Literal['x']"
is_x(reveal_type(A3.x._name_)) # N: Revealed type is "Literal['x']"
reveal_type(A3.x.value) # N: Revealed type is "builtins.int"
reveal_type(A3.x._value_) # N: Revealed type is "builtins.int"
reveal_type(A3.x.value) # N: Revealed type is "Literal[1]?"
reveal_type(A3.x._value_) # N: Revealed type is "Literal[1]?"

B1 = IntEnum('B1', 'x')
class B2(IntEnum):
Expand All @@ -639,8 +639,8 @@ reveal_type(B2.x.value) # N: Revealed type is "builtins.int"
reveal_type(B2.x._value_) # N: Revealed type is "builtins.int"
is_x(reveal_type(B3.x.name)) # N: Revealed type is "Literal['x']"
is_x(reveal_type(B3.x._name_)) # N: Revealed type is "Literal['x']"
reveal_type(B3.x.value) # N: Revealed type is "builtins.int"
reveal_type(B3.x._value_) # N: Revealed type is "builtins.int"
reveal_type(B3.x.value) # N: Revealed type is "Literal[1]?"
reveal_type(B3.x._value_) # N: Revealed type is "Literal[1]?"

# TODO: C1.x.value and C2.x.value should also be of type 'int'
# This requires either a typeshed change or a plugin refinement
Expand All @@ -661,8 +661,8 @@ reveal_type(C2.x.value) # N: Revealed type is "builtins.int"
reveal_type(C2.x._value_) # N: Revealed type is "builtins.int"
is_x(reveal_type(C3.x.name)) # N: Revealed type is "Literal['x']"
is_x(reveal_type(C3.x._name_)) # N: Revealed type is "Literal['x']"
reveal_type(C3.x.value) # N: Revealed type is "builtins.int"
reveal_type(C3.x._value_) # N: Revealed type is "builtins.int"
reveal_type(C3.x.value) # N: Revealed type is "Literal[1]?"
reveal_type(C3.x._value_) # N: Revealed type is "Literal[1]?"

D1 = Flag('D1', 'x')
class D2(Flag):
Expand All @@ -680,8 +680,8 @@ reveal_type(D2.x.value) # N: Revealed type is "builtins.int"
reveal_type(D2.x._value_) # N: Revealed type is "builtins.int"
is_x(reveal_type(D3.x.name)) # N: Revealed type is "Literal['x']"
is_x(reveal_type(D3.x._name_)) # N: Revealed type is "Literal['x']"
reveal_type(D3.x.value) # N: Revealed type is "builtins.int"
reveal_type(D3.x._value_) # N: Revealed type is "builtins.int"
reveal_type(D3.x.value) # N: Revealed type is "Literal[1]?"
reveal_type(D3.x._value_) # N: Revealed type is "Literal[1]?"

# TODO: Generalize our enum functional API logic to work with subclasses of Enum
# See https://github.com/python/mypy/issues/6037
Expand All @@ -699,8 +699,8 @@ reveal_type(E2.x.value) # N: Revealed type is "builtins.int"
reveal_type(E2.x._value_) # N: Revealed type is "builtins.int"
is_x(reveal_type(E3.x.name)) # N: Revealed type is "Literal['x']"
is_x(reveal_type(E3.x._name_)) # N: Revealed type is "Literal['x']"
reveal_type(E3.x.value) # N: Revealed type is "builtins.int"
reveal_type(E3.x._value_) # N: Revealed type is "builtins.int"
reveal_type(E3.x.value) # N: Revealed type is "Literal[1]?"
reveal_type(E3.x._value_) # N: Revealed type is "Literal[1]?"


# TODO: Figure out if we can construct enums using EnumMetas using the functional API.
Expand Down Expand Up @@ -737,9 +737,9 @@ from enum import Enum
class SomeEnum(Enum):
a = "foo"
[out]
main:2: note: Revealed type is "builtins.int"
main:2: note: Revealed type is "Literal[1]?"
[out2]
main:2: note: Revealed type is "builtins.str"
main:2: note: Revealed type is "Literal['foo']?"

[case testEnumReachabilityChecksBasic]
from enum import Enum
Expand Down Expand Up @@ -1311,8 +1311,8 @@ class Foo(Enum):
B = 2

a = Foo.A
reveal_type(a.value) # N: Revealed type is "builtins.int"
reveal_type(a._value_) # N: Revealed type is "builtins.int"
reveal_type(a.value) # N: Revealed type is "Union[Literal[1]?, Literal[2]?]"
reveal_type(a._value_) # N: Revealed type is "Union[Literal[1]?, Literal[2]?]"

[case testNewSetsUnexpectedValueType]
from enum import Enum
Expand Down Expand Up @@ -1437,7 +1437,7 @@ class Foo(Enum):
A = 1
A = 'a' # E: Attempted to reuse member name "A" in Enum definition "Foo" \
# E: Incompatible types in assignment (expression has type "str", variable has type "int")
reveal_type(Foo.A.value) # N: Revealed type is "builtins.int"
reveal_type(Foo.A.value) # N: Revealed type is "Literal[1]?"

class Bar(Enum):
A = 1
Expand Down Expand Up @@ -1796,3 +1796,18 @@ class C(Enum):
class D(C): # E: Cannot inherit from final class "C"
x: int # E: Cannot assign to final name "x"
[builtins fixtures/bool.pyi]

[case testEnumLiteralValues]
from enum import Enum

class A(Enum):
str = "foo"
int = 1
bool = False
tuple = (1,)

reveal_type(A.str.value) # N: Revealed type is "Literal['foo']?"
reveal_type(A.int.value) # N: Revealed type is "Literal[1]?"
reveal_type(A.bool.value) # N: Revealed type is "Literal[False]?"
reveal_type(A.tuple.value) # N: Revealed type is "Tuple[Literal[1]?]"
[builtins fixtures/tuple.pyi]
6 changes: 3 additions & 3 deletions test-data/unit/merge.test
Original file line number Diff line number Diff line change
Expand Up @@ -1435,16 +1435,16 @@ TypeInfo<0>(
Bases(enum.Enum<1>)
Mro(target.A<0>, enum.Enum<1>, builtins.object<2>)
Names(
X<3> (builtins.int<4>))
X<3> (Literal[0]?<4>))
MetaclassType(enum.EnumMeta<5>))
==>
TypeInfo<0>(
Name(target.A)
Bases(enum.Enum<1>)
Mro(target.A<0>, enum.Enum<1>, builtins.object<2>)
Names(
X<3> (builtins.int<4>)
Y<6> (builtins.int<4>))
X<3> (Literal[0]?<4>)
Y<6> (Literal[1]?<4>))
MetaclassType(enum.EnumMeta<5>))

[case testLiteralMerge]
Expand Down

0 comments on commit 48d810d

Please sign in to comment.