"""

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,

is_named_instance,

)

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"):

if is_named_instance(proper_type, "enum.member") and proper_type.args:

return proper_type.args[0]

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_member_callback(ctx: mypy.plugin.FunctionContext) -> Type:

"""By default `member(1)` will be infered as `member[int]`,

we want to improve the inference to be `Literal[1]` here."""

if ctx.arg_types or ctx.arg_types[0]:

arg = get_proper_type(ctx.arg_types[0][0])

proper_return = get_proper_type(ctx.default_return_type)

if (

isinstance(arg, Instance)

and arg.last_known_value

and isinstance(proper_return, Instance)

and len(proper_return.args) == 1

):

return proper_return.copy_modified(args=[arg])

return ctx.default_return_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

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, instance attributes, and `nonmember`s.

# 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 = [

_infer_value_type_with_auto_fallback(ctx, t)

for t in node_types

if t is None

or (not isinstance(t, CallableType) and not is_named_instance(t, "enum.nonmember"))

]

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