mypyc/irbuild/util.py

"""Various utilities that don't depend on other modules in mypyc.irbuild."""

from __future__ import annotations

from typing import Any, Final, Literal, TypedDict
from typing_extensions import NotRequired

from mypy.nodes import (
    ARG_NAMED,
    ARG_NAMED_OPT,
    ARG_OPT,
    ARG_POS,
    GDEF,
    ArgKind,
    AssignmentStmt,
    BytesExpr,
    CallExpr,
    ClassDef,
    Decorator,
    Expression,
    FloatExpr,
    FuncDef,
    IntExpr,
    NameExpr,
    OverloadedFuncDef,
    RefExpr,
    StrExpr,
    TempNode,
    TupleExpr,
    UnaryExpr,
    Var,
    is_class_var,
)
from mypy.semanal import refers_to_fullname
from mypy.types import FINAL_DECORATOR_NAMES
from mypyc.errors import Errors
from mypyc.ir.class_ir import ClassIR
from mypyc.ir.rtypes import RType, is_none_rprimitive, is_object_rprimitive, is_optional_type

MYPYC_ATTRS: Final[frozenset[MypycAttr]] = frozenset(
    ["native_class", "allow_interpreted_subclasses", "serializable", "free_list_len", "acyclic"]
)

DATACLASS_DECORATORS: Final = frozenset(["dataclasses.dataclass", "attr.s", "attr.attrs"])


MypycAttr = Literal[
    "native_class", "allow_interpreted_subclasses", "serializable", "free_list_len", "acyclic"
]


class MypycAttrs(TypedDict):
    native_class: NotRequired[bool]
    allow_interpreted_subclasses: NotRequired[bool]
    serializable: NotRequired[bool]
    free_list_len: NotRequired[int]
    acyclic: NotRequired[bool]


def is_final_decorator(d: Expression) -> bool:
    return refers_to_fullname(d, FINAL_DECORATOR_NAMES)


def is_trait_decorator(d: Expression) -> bool:
    return isinstance(d, RefExpr) and d.fullname == "mypy_extensions.trait"


def is_trait(cdef: ClassDef) -> bool:
    return any(is_trait_decorator(d) for d in cdef.decorators) or cdef.info.is_protocol


def dataclass_decorator_type(d: Expression) -> str | None:
    if isinstance(d, RefExpr) and d.fullname in DATACLASS_DECORATORS:
        return d.fullname.split(".")[0]
    elif (
        isinstance(d, CallExpr)
        and isinstance(d.callee, RefExpr)
        and d.callee.fullname in DATACLASS_DECORATORS
    ):
        name = d.callee.fullname.split(".")[0]
        if name == "attr" and "auto_attribs" in d.arg_names:
            # Note: the mypy attrs plugin checks that the value of auto_attribs is
            # not computed at runtime, so we don't need to perform that check here
            auto = d.args[d.arg_names.index("auto_attribs")]
            if isinstance(auto, NameExpr) and auto.name == "True":
                return "attr-auto"
        return name
    else:
        return None


def is_dataclass_decorator(d: Expression) -> bool:
    return dataclass_decorator_type(d) is not None


def is_dataclass(cdef: ClassDef) -> bool:
    return any(is_dataclass_decorator(d) for d in cdef.decorators)


# The string values returned by this function are inspected in
# mypyc/lib-rt/misc_ops.c:CPyDataclass_SleightOfHand(...).
def dataclass_type(cdef: ClassDef) -> str | None:
    for d in cdef.decorators:
        typ = dataclass_decorator_type(d)
        if typ is not None:
            return typ
    return None


def is_class_body_final(var: Var, ir: ClassIR, cdef: ClassDef, attr_rtype: RType) -> bool:
    """Is a member a final attribute initialized in class body ("X: Final = <value>")?

    The value is the same for every instance, so we don't give these struct
    slots. Instead, we let the general final-name machinery own it: the value
    is stored once in a module-level static and set on the type object (see
    ExtClassBuilder.add_attr), and reads are constant folded or loaded from
    the static.

    See ClassIR.class_final_attributes.
    """
    if not var.is_final or not var.has_explicit_value:
        return False
    if var.final_set_in_init:
        # "self.x: Final = ..." in __init__ varies per instance.
        return False
    if not ir.is_ext_class:
        # Non-extension classes keep class-level values in the class dict anyway.
        return False
    if ir.builtin_base:
        # These have no attribute slots at all (see prepare_methods_and_attributes),
        # so reads go through the class dict and honour shadowing by an interpreted
        # subclass, which folding would break.
        return False
    if dataclass_type(cdef) is not None:
        # Dataclass attribute definitions are special.
        return False
    if attr_rtype.error_overlap:
        # A static of type like this can't distinguish unset from a value equal to
        # the error sentinel.
        #
        # TODO: Support thesse by having a separate initialized flag
        return False
    return True


def _defaults_skip(stmt: AssignmentStmt, cls_type: str | None) -> bool:
    """Whether a class-level default assignment is skipped when emitting
    __mypyc_defaults_setup, based on class type.

    - attr (auto_attribs=False): skip all (handled by attr.ib machinery).
    - dataclasses / attr-auto: skip annotated assignments.
    - regular extension class: skip nothing.
    """
    if cls_type == "attr":
        return True
    if cls_type in ("dataclasses", "attr-auto"):
        return stmt.type is not None
    return False


def default_attr_name(stmt: AssignmentStmt, ir: ClassIR, cls_type: str | None) -> str | None:
    """Return the attribute name if `stmt` is a class-level default assignment
    that __mypyc_defaults_setup should emit; otherwise None.

    Single source of truth for the predicate used by both
    mypyc.irbuild.classdef.find_attr_initializers (IR build) and
    mypyc.irbuild.prepare._has_own_default_attrs (prepare-phase decl registration).
    """
    lvalue = stmt.lvalues[0]
    if not isinstance(lvalue, NameExpr) or is_class_var(lvalue):
        return None
    if isinstance(stmt.rvalue, TempNode):
        return None
    name = lvalue.name
    if name in ("__slots__", "__deletable__") or name not in ir.attributes:
        return None
    if _defaults_skip(stmt, cls_type):
        return None
    if isinstance(stmt.rvalue, RefExpr) and stmt.rvalue.fullname == "builtins.None":
        attr_type = ir.attributes[name]
        if (
            not is_optional_type(attr_type)
            and not is_object_rprimitive(attr_type)
            and not is_none_rprimitive(attr_type)
        ):
            return None
    return name


def get_mypyc_attr_literal(e: Expression) -> Any:
    """Convert an expression from a mypyc_attr decorator to a value.

    Supports a pretty limited range."""
    if isinstance(e, (StrExpr, IntExpr, FloatExpr)):
        return e.value
    elif isinstance(e, RefExpr) and e.fullname == "builtins.True":
        return True
    elif isinstance(e, RefExpr) and e.fullname == "builtins.False":
        return False
    elif isinstance(e, RefExpr) and e.fullname == "builtins.None":
        return None
    elif isinstance(e, IntExpr):
        return e.value
    return NotImplemented


def get_mypyc_attr_call(d: Expression) -> CallExpr | None:
    """Check if an expression is a call to mypyc_attr and return it if so."""
    if (
        isinstance(d, CallExpr)
        and isinstance(d.callee, RefExpr)
        and d.callee.fullname == "mypy_extensions.mypyc_attr"
    ):
        return d
    return None


def get_mypyc_attrs(
    stmt: ClassDef | Decorator, path: str, errors: Errors
) -> tuple[MypycAttrs, dict[MypycAttr, int]]:
    """Collect all the mypyc_attr attributes on a class definition or a function."""
    attrs: MypycAttrs = {}
    lines: dict[MypycAttr, int] = {}

    def set_mypyc_attr(key: str, value: Any, line: int) -> None:
        if key in MYPYC_ATTRS:
            attrs[key] = value
            lines[key] = line
        else:
            errors.error(f'"{key}" is not a supported "mypyc_attr"', path, line)
            supported_keys = '", "'.join(sorted(MYPYC_ATTRS))
            errors.note(f'supported keys: "{supported_keys}"', path, line)

    for dec in stmt.decorators:
        if d := get_mypyc_attr_call(dec):
            line = d.line
            for name, arg in zip(d.arg_names, d.args):
                if name is None:
                    if isinstance(arg, StrExpr):
                        set_mypyc_attr(arg.value, True, line)
                    else:
                        errors.error(
                            'All "mypyc_attr" positional arguments must be string literals.',
                            path,
                            line,
                        )
                else:
                    arg_value = get_mypyc_attr_literal(arg)
                    set_mypyc_attr(name, arg_value, line)

    return attrs, lines


def is_extension_class(path: str, cdef: ClassDef, errors: Errors) -> bool:
    # Check for @mypyc_attr(native_class=True/False) decorator.
    explicit_native_class = get_explicit_native_class(path, cdef, errors)

    # Classes with native_class=False are explicitly marked as non extension.
    if explicit_native_class is False:
        return False

    implicit_extension_class, reason = is_implicit_extension_class(cdef)

    # Classes with native_class=True should be extension classes, but they might
    # not be able to be due to other reasons. Print an error in that case.
    if explicit_native_class is True and not implicit_extension_class:
        errors.error(
            f"Class is marked as native_class=True but it can't be a native class. {reason}",
            path,
            cdef.line,
        )

    return implicit_extension_class


def get_explicit_native_class(path: str, cdef: ClassDef, errors: Errors) -> bool | None:
    """Return value of @mypyc_attr(native_class=True/False) decorator.

    Look for a @mypyc_attr decorator with native_class=True/False and return
    the value assigned or None if it doesn't exist. Other values are an error.
    """

    for d in cdef.decorators:
        mypyc_attr_call = get_mypyc_attr_call(d)
        if not mypyc_attr_call:
            continue

        for i, name in enumerate(mypyc_attr_call.arg_names):
            if name != "native_class":
                continue

            arg = mypyc_attr_call.args[i]
            if not isinstance(arg, NameExpr):
                errors.error("native_class must be used with True or False only", path, cdef.line)
                return None

            if arg.name == "False":
                return False
            elif arg.name == "True":
                return True
            else:
                errors.error("native_class must be used with True or False only", path, cdef.line)
                return None
    return None


def is_implicit_extension_class(cdef: ClassDef) -> tuple[bool, str]:
    """Check if class can be extension class and return a user-friendly reason it can't be one."""

    for d in cdef.decorators:
        if (
            not is_trait_decorator(d)
            and not is_dataclass_decorator(d)
            and not get_mypyc_attr_call(d)
            and not is_final_decorator(d)
        ):
            return (
                False,
                "Classes that have decorators other than supported decorators"
                " can't be native classes.",
            )

    if cdef.info.typeddict_type:
        return False, "TypedDict classes can't be native classes."
    if cdef.info.is_named_tuple:
        return False, "NamedTuple classes can't be native classes."
    if cdef.info.metaclass_type and cdef.info.metaclass_type.type.fullname not in (
        "abc.ABCMeta",
        "typing.TypingMeta",
        "typing.GenericMeta",
    ):
        return (
            False,
            "Classes with a metaclass other than ABCMeta, TypingMeta or"
            " GenericMeta can't be native classes.",
        )
    return True, ""


def get_func_def(op: FuncDef | Decorator | OverloadedFuncDef) -> FuncDef:
    if isinstance(op, OverloadedFuncDef):
        assert op.impl
        op = op.impl
    if isinstance(op, Decorator):
        op = op.func
    return op


def concrete_arg_kind(kind: ArgKind) -> ArgKind:
    """Find the concrete version of an arg kind that is being passed."""
    if kind == ARG_OPT:
        return ARG_POS
    elif kind == ARG_NAMED_OPT:
        return ARG_NAMED
    else:
        return kind


def is_constant(e: Expression) -> bool:
    """Check whether we allow an expression to appear as a default value.

    We don't currently properly support storing the evaluated
    values for default arguments and default attribute values, so
    we restrict what expressions we allow.  We allow literals of
    primitives types, None, and references to Final global
    variables.
    """
    return (
        isinstance(e, (StrExpr, BytesExpr, IntExpr, FloatExpr))
        or (isinstance(e, UnaryExpr) and e.op == "-" and isinstance(e.expr, (IntExpr, FloatExpr)))
        or (isinstance(e, TupleExpr) and all(is_constant(e) for e in e.items))
        or (
            isinstance(e, RefExpr)
            and e.kind == GDEF
            and (
                e.fullname in ("builtins.True", "builtins.False", "builtins.None")
                or (isinstance(e.node, Var) and e.node.is_final)
            )
        )
    )


def bytes_from_str(value: str) -> bytes:
    """Convert a string representing bytes into actual bytes.

    This is needed because the literal characters of BytesExpr (the
    characters inside b'') are stored in BytesExpr.value, whose type is
    'str' not 'bytes'.
    """
    return bytes(value, "utf8").decode("unicode-escape").encode("raw-unicode-escape")