Add more precise inference for enum attributes #6867

Michael0x2a · 2019-05-20T17:20:36Z

This pull request makes two changes to enum attributes.

First, this PR refines type inference for expressions like MyEnum.FOO and MyEnum.FOO.name. Those two expressions will continue to evaluate to MyEnum and str respectively under normal conditions, but will evaluate to Literal[MyEnum.FOO] and Literal["FOO"] respectively when used in Literal contexts.

Second, the type of MyEnum.FOO.value will be more precise when possible: mypy will evaluate that expression to the type of whatever FOO was assigned in the enum definition, falling back to Any as a
default.

Somewhat relatedly, this diff adds a few tests confirming we handle enum.auto() correctly.

Two additional notes:

The changes I made to the name and value fields up above are strictly speaking unsafe. While those files are normally read-only (doing MyEnum.FOO.name = blah is a runtime error), it's actually possible to change those fields anyway by altering the _name_ and _value_ fields which are not protected.

But I think this use case is probably rare -- I'm planning on investigating the feasibility of just having mypy just disallow modifying these attributes altogether after I investigate how enums are used in some internal codebases in a little more detail.
I would have liked to make MyEnum.FOO.value also return an even more precise type when used in literal contexts similar to MyEnum.FOO.name, but I think our plugin system needs to be a bit more flexible first.

This pull request makes two changes to enum attributes. First, this PR refines type inference for expressions like `MyEnum.FOO` and `MyEnum.FOO.name`. Those two expressions will continue to evaluate to `MyEnum` and `str` respectively under normal conditions, but will evaluate to `Literal[MyEnum.FOO]` and `Literal["FOO"]` respectively when used in Literal contexts. Second, the type of `MyEnum.FOO.value` will be more precise when possible: mypy will evaluate that expression to the type of whatever FOO was assigned in the enum definition, falling back to `Any` as a default. Somewhat relatedly, this diff adds a few tests confirming we handle enum.auto() correctly. Two additional notes: 1. The changes I made to the `name` and `value` fields up above are strictly speaking unsafe. While those files are normally read-only (doing `MyEnum.FOO.name = blah` is a runtime error), it's actually possible to change those fields anyway by altering the `_name_` and `_value_` fields which are *not* protected. But I think this use case is probably rare -- I'm planning on investigating the feasibility of just having mypy just disallow modifying these attributes altogether after I investigate how enums are used in some internal codebases in a little more detail. 2. I would have liked to make `MyEnum.FOO.value` also return an even more precise type when used in literal contexts similar to `MyEnum.FOO.name`, but I think our plugin system needs to be a bit more flexible first.

Michael0x2a · 2019-05-20T17:27:58Z

I guess this is kinda-sorta a follow-up to #5599?

It makes some enum attributes implicitly Final for the purposes of type inference, but doesn't disallow assignment to them. (I'm planning on tackling that in a separate PR.)

Michael0x2a · 2019-05-21T16:50:59Z

test-data/unit/check-protocols.test

@@ -2427,7 +2427,7 @@ from typing import Protocol
 class P(Protocol): ...
 class C(P): ...

-reveal_type(C.register(int))  # E: Revealed type is 'def () -> builtins.int'
+reveal_type(C.register(int))  # E: Revealed type is 'def (x: builtins.object =, base: builtins.int =) -> builtins.int'


Not sure why the revealed type is the constructor signature -- when I try directly running mypy on this, I get Type[builtins.int] instead. I guess this is just some test-related artifact?

Using a different fixture for builtins might help.

JukkaL

Looks good! It's good to have less hacky enum support. I just left some nits.

By the way, what's the expected use case for inferring a literal type for things like A.x.name? Is it to allow looking things up from a TypedDict?

What's the status of supporting Literal[Enum.foo]? The test cases don't seem to cover that, but the PR message implies that this PR is related to it.

mypy/plugins/enums.py

JukkaL · 2019-05-23T16:19:56Z

mypy/plugins/enums.py

+# Note: 'enum.EnumMeta' is deliberately excluded from this list. Classes that directly use
+# enum.EnumMeta do not necessarily automatically have the 'name' and 'value' attributes.
+ENUM_PREFIXES = ['enum.Enum', 'enum.IntEnum', 'enum.Flag', 'enum.IntFlag']
+ENUM_NAME_ACCESS = (


I wonder if using a set would be a bit faster. get_attribute_hook is called very often so it might even make a small difference.

It turns out it is indeed faster, at least based on some microbenchmarking I did. I thought the list would be small enough that overhead would be about the same either way, but that was wrong.

(In retrospect, I guess doing on average 4 to 8 some_str.__eq__(...) calls per containment check is always going to be noticeably more expensive then doing a __hash__(...) followed by maybe an __eq__(...), at least in Python.)

test-data/unit/check-enum.test

JukkaL · 2019-05-23T16:51:21Z

test-data/unit/check-protocols.test

@@ -2427,7 +2427,7 @@ from typing import Protocol
 class P(Protocol): ...
 class C(P): ...

-reveal_type(C.register(int))  # E: Revealed type is 'def () -> builtins.int'
+reveal_type(C.register(int))  # E: Revealed type is 'def (x: builtins.object =, base: builtins.int =) -> builtins.int'


Using a different fixture for builtins might help.

JukkaL · 2019-05-23T16:52:46Z

test-data/unit/lib-stub/builtins.pyi

@@ -9,6 +9,8 @@ class type:

 # These are provided here for convenience.
 class int:
+    # Note: this is a simplification of the actual signature
+    def __init__(self, x: object = ..., base: int = ...) -> None: pass


I'd rather not add anything to builtins.pyi if there's a reasonable way of avoiding it. What about adding this to, say, fixtures/primitives.pyi? How many test cases need this?

JukkaL · 2019-05-23T16:56:15Z

test-data/unit/check-enum.test

+F3.x.value     # E: "F3" has no attribute "value"
+F3.x._value_   # E: "F3" has no attribute "_value_"
+
+[case testEnumAttributeChangeIncremental]


I think that testing deserialization of the related types would also be an interesting test case. I wonder if one exists?

Hmm, I'm not sure if we have one. Do you know which file I should add the test to? (I don't remember where we keep the deserialization tests.)

mypy/plugins/enums.py

Michael0x2a · 2019-05-23T18:09:10Z

Thanks for the review! I'll work on making the changes you suggested later today.

Just to quickly answer the two questions you asked though...

I wanted to have more precise inference for things like A.x.name mostly for the sake of completeness. The general impression I got was that enums have been sort of neglected for a while in mypy and was interested in working on trying to polish them up.
Things like Literal[MyEnum.FOO] is currently supported! It snuck in here: Add basic support for enum literals #6668

JukkaL

Thanks for the updates! Just a few more issues remain.

mypy/plugins/enums.py

Michael0x2a · 2019-05-29T07:03:36Z

test-requirements.txt

@@ -4,6 +4,7 @@ flake8-bugbear; python_version >= '3.5'
 flake8-pyi; python_version >= '3.6'
 lxml==4.2.4
 mypy_extensions>=0.4.0,<0.5.0
+typing_extensions>=3.7.0,<4.0.0


I don't really understand why I needed to add this in as an explicit dependency -- or rather, how we managed to do without it up until now.

All of the other pull requests seemed fine? But I wasn't able to import typing_extensions within the new enums plugin module without it... But mypy is chock-full of imports of typing_extensions in other modules??

@Michael0x2a All the other files do:

MYPY=False if MYPY: from typing_extensions import Final

I believe you don't have a guard for the import? I think doing the same guard would fix the issue.

Argh, that would do it, thanks!

(I guess just blindly grepping for "typing_extensions" might not have quite given me the full picture...)

JukkaL

Thanks!

This pull request makes two changes to enum attributes. First, this PR refines type inference for expressions like `MyEnum.FOO` and `MyEnum.FOO.name`. Those two expressions will continue to evaluate to `MyEnum` and `str` respectively under normal conditions, but will evaluate to `Literal[MyEnum.FOO]` and `Literal["FOO"]` respectively when used in Literal contexts. Second, the type of `MyEnum.FOO.value` will be more precise when possible: mypy will evaluate that expression to the type of whatever FOO was assigned in the enum definition, falling back to `Any` as a default. Somewhat relatedly, this diff adds a few tests confirming we handle enum.auto() correctly. Two additional notes: 1. The changes I made to the `name` and `value` fields up above are strictly speaking unsafe. While those files are normally read-only (doing `MyEnum.FOO.name = blah` is a runtime error), it's actually possible to change those fields anyway by altering the `_name_` and `_value_` fields which are *not* protected. But I think this use case is probably rare -- I'm planning on investigating the feasibility of just having mypy just disallow modifying these attributes altogether after I investigate how enums are used in some internal codebases in a little more detail. 2. I would have liked to make `MyEnum.FOO.value` also return an even more precise type when used in literal contexts similar to `MyEnum.FOO.name`, but I think our plugin system needs to be a bit more flexible first.

Fix broken test

952d2d3

Michael0x2a commented May 21, 2019

View reviewed changes

JukkaL reviewed May 23, 2019

View reviewed changes

Respond to code review

e14b629

JukkaL reviewed May 28, 2019

View reviewed changes

mypy/plugins/enums.py Outdated Show resolved Hide resolved

mypy/plugins/enums.py Outdated Show resolved Hide resolved

Michael0x2a added 3 commits May 28, 2019 08:07

Respond to code review, v2

fee2827

Merge branch 'master' into enums-plugin-v2

ac20624

Experiment with adding typing_extensions as a dependency

b85ace4

Michael0x2a commented May 29, 2019

View reviewed changes

Apply Ethan's suggestion

a8d2139

JukkaL approved these changes May 30, 2019

View reviewed changes

JukkaL merged commit 39204cd into python:master May 30, 2019

saryou mentioned this pull request Mar 12, 2021

The value of enums shoulde be inferred if possible microsoft/pyright#1624

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add more precise inference for enum attributes #6867

Add more precise inference for enum attributes #6867

Michael0x2a commented May 20, 2019

Michael0x2a commented May 20, 2019

Michael0x2a May 21, 2019

JukkaL May 23, 2019

JukkaL left a comment

JukkaL May 23, 2019

Michael0x2a May 28, 2019

JukkaL May 23, 2019

JukkaL May 23, 2019

JukkaL May 23, 2019

Michael0x2a May 28, 2019

Michael0x2a commented May 23, 2019

JukkaL left a comment

Michael0x2a May 29, 2019 •

edited

ethanhs May 29, 2019

Michael0x2a May 29, 2019

JukkaL left a comment

Add more precise inference for enum attributes #6867

Add more precise inference for enum attributes #6867

Conversation

Michael0x2a commented May 20, 2019

Michael0x2a commented May 20, 2019

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JukkaL left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Michael0x2a commented May 23, 2019

JukkaL left a comment

Choose a reason for hiding this comment

Michael0x2a May 29, 2019 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JukkaL left a comment

Choose a reason for hiding this comment

Michael0x2a May 29, 2019 •

edited