259 lines
10 KiB
Python
259 lines
10 KiB
Python
|
"""
|
||
|
|
This file contains a variety of plugins for refining how mypy infers types of
|
||
|
|
expressions involving Enums.
|
||
|
|
|
||
|
|
Currently, this file focuses on providing better inference for expressions like
|
||
|
|
'SomeEnum.FOO.name' and 'SomeEnum.FOO.value'. Note that the type of both expressions
|
||
|
|
will vary depending on exactly which instance of SomeEnum we're looking at.
|
||
|
|
|
||
|
|
Note that this file does *not* contain all special-cased logic related to enums:
|
||
|
|
we actually bake some of it directly in to the semantic analysis layer (see
|
||
|
|
semanal_enum.py).
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import Final, Iterable, Sequence, TypeVar, cast
|
||
|
|
|
||
|
|
import mypy.plugin # To avoid circular imports.
|
||
|
|
from mypy.nodes import TypeInfo
|
||
|
|
from mypy.semanal_enum import ENUM_BASES
|
||
|
|
from mypy.subtypes import is_equivalent
|
||
|
|
from mypy.typeops import fixup_partial_type, make_simplified_union
|
||
|
|
from mypy.types import CallableType, Instance, LiteralType, ProperType, Type, get_proper_type
|
||
|
|
|
||
|
|
ENUM_NAME_ACCESS: Final = {f"{prefix}.name" for prefix in ENUM_BASES} | {
|
||
|
|
f"{prefix}._name_" for prefix in ENUM_BASES
|
||
|
|
}
|
||
|
|
ENUM_VALUE_ACCESS: Final = {f"{prefix}.value" for prefix in ENUM_BASES} | {
|
||
|
|
f"{prefix}._value_" for prefix in ENUM_BASES
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def enum_name_callback(ctx: mypy.plugin.AttributeContext) -> Type:
|
||
|
|
"""This plugin refines the 'name' attribute in enums to act as if
|
||
|
|
they were declared to be final.
|
||
|
|
|
||
|
|
For example, the expression 'MyEnum.FOO.name' normally is inferred
|
||
|
|
to be of type 'str'.
|
||
|
|
|
||
|
|
This plugin will instead make the inferred type be a 'str' where the
|
||
|
|
last known value is 'Literal["FOO"]'. This means it would be legal to
|
||
|
|
use 'MyEnum.FOO.name' in contexts that expect a Literal type, just like
|
||
|
|
any other Final variable or attribute.
|
||
|
|
|
||
|
|
This plugin assumes that the provided context is an attribute access
|
||
|
|
matching one of the strings found in 'ENUM_NAME_ACCESS'.
|
||
|
|
"""
|
||
|
|
enum_field_name = _extract_underlying_field_name(ctx.type)
|
||
|
|
if enum_field_name is None:
|
||
|
|
return ctx.default_attr_type
|
||
|
|
else:
|
||
|
|
str_type = ctx.api.named_generic_type("builtins.str", [])
|
||
|
|
literal_type = LiteralType(enum_field_name, fallback=str_type)
|
||
|
|
return str_type.copy_modified(last_known_value=literal_type)
|
||
|
|
|
||
|
|
|
||
|
|
_T = TypeVar("_T")
|
||
|
|
|
||
|
|
|
||
|
|
def _first(it: Iterable[_T]) -> _T | None:
|
||
|
|
"""Return the first value from any iterable.
|
||
|
|
|
||
|
|
Returns ``None`` if the iterable is empty.
|
||
|
|
"""
|
||
|
|
for val in it:
|
||
|
|
return val
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _infer_value_type_with_auto_fallback(
|
||
|
|
ctx: mypy.plugin.AttributeContext, proper_type: ProperType | None
|
||
|
|
) -> Type | None:
|
||
|
|
"""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
|
||
|
|
proper_type = get_proper_type(fixup_partial_type(proper_type))
|
||
|
|
if not (isinstance(proper_type, Instance) and proper_type.type.fullname == "enum.auto"):
|
||
|
|
return proper_type
|
||
|
|
assert isinstance(ctx.type, Instance), "An incorrect ctx.type was passed."
|
||
|
|
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_gnv = _first(ti for ti in info.mro if ti.names.get("_generate_next_value_"))
|
||
|
|
if type_with_gnv is None:
|
||
|
|
return ctx.default_attr_type
|
||
|
|
|
||
|
|
stnode = type_with_gnv.names["_generate_next_value_"]
|
||
|
|
|
||
|
|
# This should be a `CallableType`
|
||
|
|
node_type = get_proper_type(stnode.type)
|
||
|
|
if isinstance(node_type, CallableType):
|
||
|
|
if type_with_gnv.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
|
||
|
|
|
||
|
|
|
||
|
|
def _implements_new(info: TypeInfo) -> bool:
|
||
|
|
"""Check whether __new__ comes from enum.Enum or was implemented in a
|
||
|
|
subclass. In the latter case, we must infer Any as long as mypy can't infer
|
||
|
|
the type of _value_ from assignments in __new__.
|
||
|
|
"""
|
||
|
|
type_with_new = _first(
|
||
|
|
ti
|
||
|
|
for ti in info.mro
|
||
|
|
if ti.names.get("__new__") and not ti.fullname.startswith("builtins.")
|
||
|
|
)
|
||
|
|
if type_with_new is None:
|
||
|
|
return False
|
||
|
|
return type_with_new.fullname not in ("enum.Enum", "enum.IntEnum", "enum.StrEnum")
|
||
|
|
|
||
|
|
|
||
|
|
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
|
||
|
|
following:
|
||
|
|
|
||
|
|
class SomeEnum:
|
||
|
|
FOO = A()
|
||
|
|
BAR = B()
|
||
|
|
|
||
|
|
By default, mypy will infer that 'SomeEnum.FOO.value' and
|
||
|
|
'SomeEnum.BAR.value' both are of type 'Any'. This plugin refines
|
||
|
|
this inference so that mypy understands the expressions are
|
||
|
|
actually of types 'A' and 'B' respectively. This better reflects
|
||
|
|
the actual runtime behavior.
|
||
|
|
|
||
|
|
This plugin works simply by looking up the original value assigned
|
||
|
|
to the enum. For example, when this plugin sees 'SomeEnum.BAR.value',
|
||
|
|
it will look up whatever type 'BAR' had in the SomeEnum TypeInfo and
|
||
|
|
use that as the inferred type of the overall expression.
|
||
|
|
|
||
|
|
This plugin assumes that the provided context is an attribute access
|
||
|
|
matching one of the strings found in 'ENUM_VALUE_ACCESS'.
|
||
|
|
"""
|
||
|
|
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
|
||
|
|
|
||
|
|
# As long as mypy doesn't understand attribute creation in __new__,
|
||
|
|
# there is no way to predict the value type if the enum class has a
|
||
|
|
# custom implementation
|
||
|
|
if _implements_new(info):
|
||
|
|
return ctx.default_attr_type
|
||
|
|
|
||
|
|
stnodes = (info.get(name) for name in info.names)
|
||
|
|
|
||
|
|
# Enums _can_ have methods and instance attributes.
|
||
|
|
# Omit methods and attributes created by assigning to self.*
|
||
|
|
# for our value inference.
|
||
|
|
node_types = (
|
||
|
|
get_proper_type(n.type) if n else None
|
||
|
|
for n in stnodes
|
||
|
|
if n is None or not n.implicit
|
||
|
|
)
|
||
|
|
proper_types = list(
|
||
|
|
_infer_value_type_with_auto_fallback(ctx, t)
|
||
|
|
for t in node_types
|
||
|
|
if t is None or not isinstance(t, CallableType)
|
||
|
|
)
|
||
|
|
underlying_type = _first(proper_types)
|
||
|
|
if underlying_type is None:
|
||
|
|
return ctx.default_attr_type
|
||
|
|
|
||
|
|
# At first we try to predict future `value` type if all other items
|
||
|
|
# have the same type. For example, `int`.
|
||
|
|
# If this is the case, we simply return this type.
|
||
|
|
# See https://github.com/python/mypy/pull/9443
|
||
|
|
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
|
||
|
|
|
||
|
|
# But, after we started treating all `Enum` values as `Final`,
|
||
|
|
# we start to infer types in
|
||
|
|
# `item = 1` as `Literal[1]`, not just `int`.
|
||
|
|
# So, for example types in this `Enum` will all be different:
|
||
|
|
#
|
||
|
|
# class Ordering(IntEnum):
|
||
|
|
# one = 1
|
||
|
|
# two = 2
|
||
|
|
# three = 3
|
||
|
|
#
|
||
|
|
# We will infer three `Literal` types here.
|
||
|
|
# They are not the same, but they are equivalent.
|
||
|
|
# So, we unify them to make sure `.value` prediction still works.
|
||
|
|
# Result will be `Literal[1] | Literal[2] | Literal[3]` for this case.
|
||
|
|
all_equivalent_types = all(
|
||
|
|
proper_type is not None and is_equivalent(proper_type, underlying_type)
|
||
|
|
for proper_type in proper_types
|
||
|
|
)
|
||
|
|
if all_equivalent_types:
|
||
|
|
return make_simplified_union(cast(Sequence[Type], proper_types))
|
||
|
|
return ctx.default_attr_type
|
||
|
|
|
||
|
|
assert isinstance(ctx.type, Instance)
|
||
|
|
info = ctx.type.type
|
||
|
|
|
||
|
|
# As long as mypy doesn't understand attribute creation in __new__,
|
||
|
|
# there is no way to predict the value type if the enum class has a
|
||
|
|
# custom implementation
|
||
|
|
if _implements_new(info):
|
||
|
|
return ctx.default_attr_type
|
||
|
|
|
||
|
|
stnode = info.get(enum_field_name)
|
||
|
|
if stnode is None:
|
||
|
|
return ctx.default_attr_type
|
||
|
|
|
||
|
|
underlying_type = _infer_value_type_with_auto_fallback(ctx, get_proper_type(stnode.type))
|
||
|
|
if underlying_type is None:
|
||
|
|
return ctx.default_attr_type
|
||
|
|
|
||
|
|
return underlying_type
|
||
|
|
|
||
|
|
|
||
|
|
def _extract_underlying_field_name(typ: Type) -> str | None:
|
||
|
|
"""If the given type corresponds to some Enum instance, returns the
|
||
|
|
original name of that enum. For example, if we receive in the type
|
||
|
|
corresponding to 'SomeEnum.FOO', we return the string "SomeEnum.Foo".
|
||
|
|
|
||
|
|
This helper takes advantage of the fact that Enum instances are valid
|
||
|
|
to use inside Literal[...] types. An expression like 'SomeEnum.FOO' is
|
||
|
|
actually represented by an Instance type with a Literal enum fallback.
|
||
|
|
|
||
|
|
We can examine this Literal fallback to retrieve the string.
|
||
|
|
"""
|
||
|
|
typ = get_proper_type(typ)
|
||
|
|
if not isinstance(typ, Instance):
|
||
|
|
return None
|
||
|
|
|
||
|
|
if not typ.type.is_enum:
|
||
|
|
return None
|
||
|
|
|
||
|
|
underlying_literal = typ.last_known_value
|
||
|
|
if underlying_literal is None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
# The checks above have verified this LiteralType is representing an enum value,
|
||
|
|
# which means the 'value' field is guaranteed to be the name of the enum field
|
||
|
|
# as a string.
|
||
|
|
assert isinstance(underlying_literal.value, str)
|
||
|
|
return underlying_literal.value
|