"""Semantic analysis of call-based Enum definitions.

This is conceptually part of mypy.semanal (semantic analyzer pass 2).

"""

from __future__ import annotations

from typing import Final, cast

from mypy.nodes import (

ARG_NAMED,

ARG_POS,

MDEF,

AssignmentStmt,

CallExpr,

Context,

DictExpr,

EnumCallExpr,

Expression,

ListExpr,

MemberExpr,

NameExpr,

RefExpr,

StrExpr,

SymbolTableNode,

TupleExpr,

TypeInfo,

Var,

is_StrExpr_list,

)

from mypy.options import Options

from mypy.semanal_shared import SemanticAnalyzerInterface

from mypy.types import ENUM_REMOVED_PROPS, LiteralType, 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.

ENUM_BASES: Final = frozenset(

("enum.Enum", "enum.IntEnum", "enum.Flag", "enum.IntFlag", "enum.StrEnum")

)

ENUM_SPECIAL_PROPS: Final = frozenset(

(

"name",

"value",

"_name_",

"_value_",

*ENUM_REMOVED_PROPS,

# Also attributes from `object`:

"__module__",

"__annotations__",

"__doc__",

"__slots__",

"__dict__",

)

)

class EnumCallAnalyzer:

def __init__(self, options: Options, api: SemanticAnalyzerInterface) -> None:

self.options = options

self.api = api

def process_enum_call(self, s: AssignmentStmt, is_func_scope: bool) -> bool:

"""Check if s defines an Enum; if yes, store the definition in symbol table.

Return True if this looks like an Enum definition (but maybe with errors),

otherwise return False.

"""

if len(s.lvalues) != 1 or not isinstance(s.lvalues[0], (NameExpr, MemberExpr)):

return False

lvalue = s.lvalues[0]

name = lvalue.name

enum_call = self.check_enum_call(s.rvalue, name, is_func_scope)

if enum_call is None:

return False

if isinstance(lvalue, MemberExpr):

self.fail("Enum type as attribute is not supported", lvalue)

return False

# Yes, it's a valid Enum definition. Add it to the symbol table.

self.api.add_symbol(name, enum_call, s)

return True

def check_enum_call(

self, node: Expression, var_name: str, is_func_scope: bool

) -> TypeInfo | None:

"""Check if a call defines an Enum.

Example:

A = enum.Enum('A', 'foo bar')

is equivalent to:

class A(enum.Enum):

foo = 1

bar = 2

"""

if not isinstance(node, CallExpr):

return None

call = node

callee = call.callee

if not isinstance(callee, RefExpr):

return None

fullname = callee.fullname

if fullname not in ENUM_BASES:

return None

new_class_name, items, values, ok = self.parse_enum_call_args(

call, fullname.split(".")[-1]

)

if not ok:

# Error. Construct dummy return value.

name = var_name

if is_func_scope:

name += "@" + str(call.line)

info = self.build_enum_call_typeinfo(name, [], fullname, node.line)

else:

if new_class_name != var_name:

msg = f'String argument 1 "{new_class_name}" to {fullname}(...) does not match variable name "{var_name}"'

self.fail(msg, call)

name = cast(StrExpr, call.args[0]).value

if name != var_name or is_func_scope:

# Give it a unique name derived from the line number.

name += "@" + str(call.line)

info = self.build_enum_call_typeinfo(name, items, fullname, call.line)

# Store generated TypeInfo under both names, see semanal_namedtuple for more details.

if name != var_name or is_func_scope:

self.api.add_symbol_skip_local(name, info)

call.analyzed = EnumCallExpr(info, items, values)

call.analyzed.set_line(call)

info.line = node.line

return info

def build_enum_call_typeinfo(

self, name: str, items: list[str], fullname: str, line: int

) -> TypeInfo:

base = self.api.named_type_or_none(fullname)

assert base is not None

info = self.api.basic_new_typeinfo(name, base, line)

info.metaclass_type = info.calculate_metaclass_type()

info.is_enum = True

for item in items:

var = Var(item)

var.info = info

var.is_property = True

var._fullname = f"{info.fullname}.{item}"

info.names[item] = SymbolTableNode(MDEF, var)

return info

def parse_enum_call_args(

self, call: CallExpr, class_name: str

) -> tuple[str, list[str], list[Expression | None], bool]:

"""Parse arguments of an Enum call.

Return a tuple of fields, values, was there an error.

"""

args = call.args

if not all(arg_kind in [ARG_POS, ARG_NAMED] for arg_kind in call.arg_kinds):

return self.fail_enum_call_arg(f"Unexpected arguments to {class_name}()", call)

if len(args) < 2:

return self.fail_enum_call_arg(f"Too few arguments for {class_name}()", call)

if len(args) > 6:

return self.fail_enum_call_arg(f"Too many arguments for {class_name}()", call)

valid_name = [None, "value", "names", "module", "qualname", "type", "start"]

for arg_name in call.arg_names:

if arg_name not in valid_name:

self.fail_enum_call_arg(f'Unexpected keyword argument "{arg_name}"', call)

value, names = None, None

for arg_name, arg in zip(call.arg_names, args):

if arg_name == "value":

value = arg

if arg_name == "names":

names = arg

if value is None:

value = args[0]

if names is None:

names = args[1]

if not isinstance(value, StrExpr):

return self.fail_enum_call_arg(

f"{class_name}() expects a string literal as the first argument", call

)

new_class_name = value.value

items = []

values: list[Expression | None] = []

if isinstance(names, StrExpr):

fields = names.value

for field in fields.replace(",", " ").split():

items.append(field)

elif isinstance(names, (TupleExpr, ListExpr)):

seq_items = names.items

if is_StrExpr_list(seq_items):

items = [seq_item.value for seq_item in seq_items]

elif all(

isinstance(seq_item, (TupleExpr, ListExpr))

and len(seq_item.items) == 2

and isinstance(seq_item.items[0], StrExpr)

for seq_item in seq_items

):

for seq_item in seq_items:

assert isinstance(seq_item, (TupleExpr, ListExpr))

name, value = seq_item.items

assert isinstance(name, StrExpr)

items.append(name.value)

values.append(value)

else:

return self.fail_enum_call_arg(

"%s() with tuple or list expects strings or (name, value) pairs" % class_name,

call,

)

elif isinstance(names, DictExpr):

for key, value in names.items:

if not isinstance(key, StrExpr):

return self.fail_enum_call_arg(

f"{class_name}() with dict literal requires string literals", call

)

items.append(key.value)

values.append(value)

elif isinstance(args[1], RefExpr) and isinstance(args[1].node, Var):

proper_type = get_proper_type(args[1].node.type)

if (

proper_type is not None

and isinstance(proper_type, LiteralType)

and isinstance(proper_type.value, str)

):

fields = proper_type.value

for field in fields.replace(",", " ").split():

items.append(field)

elif args[1].node.is_final and isinstance(args[1].node.final_value, str):

fields = args[1].node.final_value

for field in fields.replace(",", " ").split():

items.append(field)

else:

return self.fail_enum_call_arg(

"Second argument of %s() must be string, tuple, list or dict literal for mypy to determine Enum members"

% class_name,

call,

)

else:

# TODO: Allow dict(x=1, y=2) as a substitute for {'x': 1, 'y': 2}?

return self.fail_enum_call_arg(

"Second argument of %s() must be string, tuple, list or dict literal for mypy to determine Enum members"

% class_name,

call,

)

if not items:

return self.fail_enum_call_arg(f"{class_name}() needs at least one item", call)

if not values:

values = [None] * len(items)

assert len(items) == len(values)

return new_class_name, items, values, True

def fail_enum_call_arg(

self, message: str, context: Context

) -> tuple[str, list[str], list[Expression | None], bool]:

self.fail(message, context)

return "", [], [], False

# Helpers

def fail(self, msg: str, ctx: Context) -> None:

self.api.fail(msg, ctx)