Skip to content

Predict enum value type for unknown member names. #9443

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.

Sign up for GitHub

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

Merged
gvanrossum merged 13 commits into from
Sep 17, 2020
Merged

Predict enum value type for unknown member names. #9443

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
fa67a31
Predict enum value type for unknown member names.
mgilson-argo Sep 15, 2020
077d3e7
Fix a typo. Squash before merging.
mgilson-argo Sep 15, 2020
e861117
Get enum-value inference to work with auto.
mgilson-argo Sep 16, 2020
db16871
Fix flake8 complaint.
mgilson-argo Sep 16, 2020
bfee76d
Fixup the unit-tests.
mgilson-argo Sep 16, 2020
b3653d7
Remove some unused fixtures.
mgilson-argo Sep 16, 2020
f1adaf5
Make the self-test happy.
mgilson-argo Sep 16, 2020
22c1474
Invert logic for inferring with value-type in presence of auto.
mgilson-argo Sep 17, 2020
aae9de3
Change a ValueError to an `assert` for type collapsing.
mgilson-argo Sep 17, 2020
d392d05
Shorten some names and drop a branch or two.
mgilson-argo Sep 17, 2020
050fdc2
Get my conditional right ... I think ...
mgilson-argo Sep 17, 2020
5fa5a93
Fix the fixture application.
mgilson-argo Sep 17, 2020
5b8baa9
Stop using `next` for now.
mgilson-argo Sep 17, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 71 additions & 9 deletions mypy/plugins/enums.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from typing_extensions import Final

import mypy.plugin # To avoid circular imports.
from mypy.types import Type, Instance, LiteralType, get_proper_type
from mypy.types import Type, Instance, LiteralType, CallableType, ProperType, get_proper_type

# 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.
Expand Down Expand Up @@ -53,6 +53,48 @@ def enum_name_callback(ctx: 'mypy.plugin.AttributeContext') -> Type:
return str_type.copy_modified(last_known_value=literal_type)


def _infer_value_type_with_auto_fallback(
ctx: 'mypy.plugin.AttributeContext',
proper_type: Optional[ProperType]) -> Optional[Type]:
"""Figure out the type of an enum value accounting for `auto()`.

This method is a no-op for a `None` proper_type and also in the case where
the type is not "enum.auto"
"""
if proper_type is None:
return None
if (isinstance(proper_type, Instance) and
proper_type.type.fullname == 'enum.auto'):
Comment thread
gvanrossum marked this conversation as resolved.
Outdated
if not isinstance(ctx.type, Instance):
raise ValueError("An incorrect ctx.type was passed.")
Comment thread
gvanrossum marked this conversation as resolved.
Outdated
info = ctx.type.type
# Find the first _generate_next_value_ on the mro. We need to know
# if it is `Enum` because `Enum` types say that the return-value of
# `_generate_next_value_` is `Any`. In reality the default `auto()`
# returns an `int` (presumably the `Any` in typeshed is to make it
# easier to subclass and change the returned type).
type_with_generate_next_value = next(
Comment thread
gvanrossum marked this conversation as resolved.
Outdated
(type_info for type_info in info.mro
if type_info.names.get('_generate_next_value_')),
None)
if type_with_generate_next_value is None:
return ctx.default_attr_type

stnode = type_with_generate_next_value.get('_generate_next_value_')
if stnode is None:
return ctx.default_attr_type
Comment thread
gvanrossum marked this conversation as resolved.
Outdated

# This should be a `CallableType`
node_type = get_proper_type(stnode.type)
if isinstance(node_type, CallableType):
if type_with_generate_next_value.fullname == 'enum.Enum':
int_type = ctx.api.named_generic_type('builtins.int', [])
return int_type
return get_proper_type(node_type.ret_type)
return ctx.default_attr_type
return proper_type


def enum_value_callback(ctx: 'mypy.plugin.AttributeContext') -> Type:
"""This plugin refines the 'value' attribute in enums to refer to
the original underlying value. For example, suppose we have the
Expand All @@ -78,6 +120,32 @@ class SomeEnum:
"""
enum_field_name = _extract_underlying_field_name(ctx.type)
if enum_field_name is None:
# We do not know the enum field name (perhaps it was passed to a
# function and we only know that it _is_ a member). All is not lost
# however, if we can prove that the all of the enum members have the
# same value-type, then it doesn't matter which member was passed in.
# The value-type is still known.
if isinstance(ctx.type, Instance):
info = ctx.type.type
Comment thread
gvanrossum marked this conversation as resolved.
stnodes = (info.get(name) for name in info.names)
# Enums _can_ have methods.
# Omit methods for our value inference.
node_types = (
get_proper_type(n.type) if n else None
for n in stnodes)
proper_types = (
_infer_value_type_with_auto_fallback(ctx, t)
for t in node_types
if t is None or not isinstance(t, CallableType))
underlying_type = next(proper_types, None)
if underlying_type is None:
return ctx.default_attr_type
all_same_value_type = all(
proper_type is not None and proper_type == underlying_type
for proper_type in proper_types)
if all_same_value_type:
if underlying_type is not None:
return underlying_type

@gvanrossum gvanrossum Sep 15, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I read the docstring of get_proper_type() correctly, it suggest that here you should actually be returning first_node_type. And in fact, perhaps you should probably adjust the all() check to go through get_proper_type() instead?

@mgilson mgilson Sep 16, 2020

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I took a stab at this, but I'm not completely sure what you're saying here :). It's unclear to me whether the result of enum_value_callback should be a ProperType or not. The existing code seems to try to return a ProperType so I just followed suit there. I did change the "all-equal" to compare the entities after figuring out their ProperType though.

@gvanrossum gvanrossum Sep 17, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hm, I think I misunderstood this. I thought that using get_proper_type() would distinguish A from int if we have

A = int
def f(a: A): ...

but it doesn't seem to work that way. So you're forgiven for not understanding me. :-)

I think what you have now is fine.

return ctx.default_attr_type

assert isinstance(ctx.type, Instance)
Comment thread
gvanrossum marked this conversation as resolved.
Expand All @@ -86,15 +154,9 @@ class SomeEnum:
if stnode is None:
return ctx.default_attr_type

underlying_type = get_proper_type(stnode.type)
underlying_type = _infer_value_type_with_auto_fallback(
ctx, get_proper_type(stnode.type))
if underlying_type is None:
# TODO: Deduce the inferred type if the user omits adding their own default types.
# TODO: Consider using the return type of `Enum._generate_next_value_` here?
return ctx.default_attr_type

if isinstance(underlying_type, Instance) and underlying_type.type.fullname == 'enum.auto':
# TODO: Deduce the correct inferred type when the user uses 'enum.auto'.
# We should use the same strategy we end up picking up above.
return ctx.default_attr_type

return underlying_type
Expand Down
86 changes: 77 additions & 9 deletions test-data/unit/check-enum.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,74 @@ reveal_type(Truth.true.name) # N: Revealed type is 'Literal['true']?'
reveal_type(Truth.false.value) # N: Revealed type is 'builtins.bool'
[builtins fixtures/bool.pyi]

[case testEnumValueExtended]
from enum import Enum
class Truth(Enum):
true = True
false = False

def infer_truth(truth: Truth) -> None:
reveal_type(truth.value) # N: Revealed type is 'builtins.bool'
[builtins fixtures/bool.pyi]

[case testEnumValueAllAuto]
from enum import Enum, auto
class Truth(Enum):
true = auto()
false = auto()

def infer_truth(truth: Truth) -> None:
reveal_type(truth.value) # N: Revealed type is 'builtins.int'
[builtins fixtures/primitives.pyi]

[case testEnumValueSomeAuto]
from enum import Enum, auto
class Truth(Enum):
true = 8675309
false = auto()

def infer_truth(truth: Truth) -> None:
reveal_type(truth.value) # N: Revealed type is 'builtins.int'

[case testEnumValueExtraMethods]
from enum import Enum, auto
class Truth(Enum):
true = True
false = False

def foo(self) -> str:
return 'bar'

def infer_truth(truth: Truth) -> None:
reveal_type(truth.value) # N: Revealed type is 'builtins.bool'
[builtins fixtures/bool.pyi]

[case testEnumValueCustomAuto]
from enum import Enum, auto
class AutoName(Enum):

# In `typeshed`, this is a staticmethod and has more arguments,
# but I have lied a bit to keep the test stubs lean.
def _generate_next_value_(self) -> str:
return "name"

class Truth(AutoName):
true = auto()
false = auto()

def infer_truth(truth: Truth) -> None:
reveal_type(truth.value) # N: Revealed type is 'builtins.str'

[case testEnumValueInhomogenous]
from enum import Enum
class Truth(Enum):
true = 'True'
false = 0

def cannot_infer_truth(truth: Truth) -> None:
reveal_type(truth.value) # N: Revealed type is 'Any'

@gvanrossum gvanrossum Sep 15, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if this shouldn't infer object instead of Any?

In favor of Any: That's what we did before in this case; it prevents false positives.

In favor of object: That's what you'd see for e.g. reveal_type("" if a else 0), assuming a has an unknown value.

@mgilson mgilson Sep 16, 2020

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was trying to just follow the existing art. I imagine that changing to object could cause some code to fail with false positives (If a user is actually using the .value, it seems unlikely that they will be restricting themselves to the interface provided by object so this would probably force them to do additional cast or other methods of type-narrowing)

With that said ... once this is implemented, it shouldn't matter to me either way. Hopefully mypy will know the types of all of my enum values since I see no reason for inhomogenous values and it should never bother me in either case :). I'll do whatever you like.

@gvanrossum gvanrossum Sep 17, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, let's stick with tradition and keep Any.

[builtins fixtures/bool.pyi]

[case testEnumUnique]
import enum
@enum.unique
Expand Down Expand Up @@ -497,8 +565,8 @@ reveal_type(A1.x.value) # N: Revealed type is 'Any'
reveal_type(A1.x._value_) # N: Revealed type is 'Any'
is_x(reveal_type(A2.x.name)) # N: Revealed type is 'Literal['x']'
is_x(reveal_type(A2.x._name_)) # N: Revealed type is 'Literal['x']'
reveal_type(A2.x.value) # N: Revealed type is 'Any'
reveal_type(A2.x._value_) # N: Revealed type is 'Any'
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'
Expand All @@ -519,7 +587,7 @@ reveal_type(B1.x._value_) # N: Revealed type is 'Any'
is_x(reveal_type(B2.x.name)) # N: Revealed type is 'Literal['x']'
is_x(reveal_type(B2.x._name_)) # N: Revealed type is 'Literal['x']'
reveal_type(B2.x.value) # N: Revealed type is 'builtins.int'
reveal_type(B2.x._value_) # N: Revealed type is 'Any'
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'
Expand All @@ -540,8 +608,8 @@ reveal_type(C1.x.value) # N: Revealed type is 'Any'
reveal_type(C1.x._value_) # N: Revealed type is 'Any'
is_x(reveal_type(C2.x.name)) # N: Revealed type is 'Literal['x']'
is_x(reveal_type(C2.x._name_)) # N: Revealed type is 'Literal['x']'
reveal_type(C2.x.value) # N: Revealed type is 'Any'
reveal_type(C2.x._value_) # N: Revealed type is 'Any'
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'
Expand All @@ -559,8 +627,8 @@ reveal_type(D1.x.value) # N: Revealed type is 'Any'
reveal_type(D1.x._value_) # N: Revealed type is 'Any'
is_x(reveal_type(D2.x.name)) # N: Revealed type is 'Literal['x']'
is_x(reveal_type(D2.x._name_)) # N: Revealed type is 'Literal['x']'
reveal_type(D2.x.value) # N: Revealed type is 'Any'
reveal_type(D2.x._value_) # N: Revealed type is 'Any'
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'
Expand All @@ -578,8 +646,8 @@ class E3(Parent):

is_x(reveal_type(E2.x.name)) # N: Revealed type is 'Literal['x']'
is_x(reveal_type(E2.x._name_)) # N: Revealed type is 'Literal['x']'
reveal_type(E2.x.value) # N: Revealed type is 'Any'
reveal_type(E2.x._value_) # N: Revealed type is 'Any'
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'
Expand Down
9 changes: 8 additions & 1 deletion test-data/unit/lib-stub/enum.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@ class Enum(metaclass=EnumMeta):
_name_: str
_value_: Any

# In reality, _generate_next_value_ is python3.6 only and has a different signature.

@mgilson mgilson Sep 16, 2020

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shenanigans are happening here -- I'm not sure the best way to do this. If I use the signature (as written in typeshed), we end up with:

        @staticmethod
        def _generate_next_value_(name: str, start: int, count: int, last_values: List[Any]) -> Any: pass

That pulls in staticmethod which doesn't have a stub in lib-stub. From that point, there are a bunch of different ways we can lie (e.g. just add a self and drop the staticmethod), but it seems to me like a little lie is just as harmful as a big one (maybe more because it's less likely to be noticed).

@gvanrossum gvanrossum Sep 17, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this comment is all we need if in the future mypy and/or the test fixtures somehow get clever enough to notice the signature mismatch.

# However, this should be quick and doesn't require additional stubs (e.g. `staticmethod`)
def _generate_next_value_(self) -> Any: pass

class IntEnum(int, Enum):
value: int

Expand All @@ -37,4 +41,7 @@ class IntFlag(int, Flag):


class auto(IntFlag):
value: Any

value: Any

def __init__(self) -> None: pass

@mgilson mgilson Sep 16, 2020

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know why I need to add this. It isn't present in typeshed, but my actual runs still seem to work. Without it, I have two tests that fail with complaints that auto didn't get the right number of arguments:

[case testEnumValueSomeAuto]
from enum import Enum, auto
class Truth(Enum):
    true = 8675309
    false = auto()

def infer_truth(truth: Truth) -> None:
    reveal_type(truth.value) # N: Revealed type is 'builtins.int'

And

[case testEnumValueCustomAuto]
from enum import Enum, auto
from typing import List, Any
class AutoName(Enum):

    # In `typeshed`, this is a staticmethod and has more arguments,
    # but I have lied a bit to keep the test stubs lean.
    def _generate_next_value_(self) -> str:
        return "name"

class Truth(AutoName):
    true = auto()
    false = auto()

def infer_truth(truth: Truth) -> None:
    reveal_type(truth.value) # N: Revealed type is 'builtins.str'

If anyone can help me reason through why that would be, I'm happy to dig deeper.

@gvanrossum gvanrossum Sep 16, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's because auto inherits Enum.__new__ and that has a mandatory argument.

@mgilson mgilson Sep 17, 2020

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had the same thought, but as far as I can tell, the same thing happens in the official stubs in typeshed -- but we don't see this complaint when I run it for real. It also only complains for these two tests (which admittedly are doing something a little special with auto) -- even though there are a lot of other tests that use auto. I've never really seen mypy come up with a different required signature for something based on surrounding context. Maybe there is something interesting happening in the plugin that I haven't quite grokked yet -- particularly when mixed with the slightly less-than-usual stub setup for these tests?

@gvanrossum gvanrossum Sep 17, 2020

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding [builtins fixtures/primitives.pyi] to those two tests resolves this. I think I have an idea: without that, you get int and object from the very basic lib-stub/builtins.pyi, and there int has no __init__. But in fixtures/primitives.pyi there is an int.__init__.

The secret is in the MRO of auto, which is (in either case) auto, IntFlag, int, Flag, Enum, object. With primitives.pyi, we pick up __init__ from int. But without it, we pick it up from Enum... And Enum.__init__ requires an argument.

So it's purely a test issue, and IMO the fix is to add [builtins fixtures/primitives.pyi] and get rid of auto.__init__. (And of the blank line before value: Any.)