Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add more precise inference for enum attributes #6867

Merged
merged 7 commits into from May 30, 2019

Conversation

@Michael0x2a
Copy link
Collaborator

commented May 20, 2019

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.

Add more precise inference for enum attributes
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

This comment has been minimized.

Copy link
Collaborator Author

commented May 20, 2019

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.)

@@ -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'

This comment has been minimized.

Copy link
@Michael0x2a

Michael0x2a May 21, 2019

Author Collaborator

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?

This comment has been minimized.

Copy link
@JukkaL

JukkaL May 23, 2019

Collaborator

Using a different fixture for builtins might help.

@JukkaL
Copy link
Collaborator

left a comment

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 Outdated Show resolved Hide resolved
# 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 = (

This comment has been minimized.

Copy link
@JukkaL

JukkaL May 23, 2019

Collaborator

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.

This comment has been minimized.

Copy link
@Michael0x2a

Michael0x2a May 28, 2019

Author Collaborator

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 Outdated Show resolved Hide resolved
@@ -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'

This comment has been minimized.

Copy link
@JukkaL

JukkaL May 23, 2019

Collaborator

Using a different fixture for builtins might help.

@@ -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

This comment has been minimized.

Copy link
@JukkaL

JukkaL May 23, 2019

Collaborator

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?

F3.x.value # E: "F3" has no attribute "value"
F3.x._value_ # E: "F3" has no attribute "_value_"

[case testEnumAttributeChangeIncremental]

This comment has been minimized.

Copy link
@JukkaL

JukkaL May 23, 2019

Collaborator

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

This comment has been minimized.

Copy link
@Michael0x2a

Michael0x2a May 28, 2019

Author Collaborator

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 Show resolved Hide resolved
mypy/plugins/enums.py Show resolved Hide resolved
mypy/plugins/enums.py Show resolved Hide resolved
@Michael0x2a

This comment has been minimized.

Copy link
Collaborator Author

commented May 23, 2019

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...

  1. 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.

  2. Things like Literal[MyEnum.FOO] is currently supported! It snuck in here: #6668

@JukkaL
Copy link
Collaborator

left a comment

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

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

Michael0x2a added some commits May 28, 2019

@@ -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

This comment has been minimized.

Copy link
@Michael0x2a

Michael0x2a May 29, 2019

Author Collaborator

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??

This comment has been minimized.

Copy link
@ethanhs

ethanhs May 29, 2019

Collaborator

@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.

This comment has been minimized.

Copy link
@Michael0x2a

Michael0x2a May 29, 2019

Author Collaborator

Argh, that would do it, thanks!

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

@JukkaL

JukkaL approved these changes May 30, 2019

Copy link
Collaborator

left a comment

Thanks!

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

2 checks passed

continuous-integration/appveyor/pr AppVeyor build succeeded
Details
continuous-integration/travis-ci/pr The Travis CI build passed
Details

PattenR added a commit to PattenR/mypy that referenced this pull request Jun 23, 2019

Add more precise inference for enum attributes (python#6867)
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
3 participants
You can’t perform that action at this time.