Class EnumCallAnalyzer

mypy/semanal_enum.py:58–269

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)