Attributes

def param_def(
    constraint: "IRDLAttrConstraint[AttributeInvT] | None" = None,
    *,
    converter: Callable[[Any], AttributeInvT] | None = None,
    init: Literal[True] = True,
) -> AttributeInvT:
    """Defines a parameter of a ParametrizedAttribute."""
    return cast(
        AttributeInvT,
        ParamDef(
            AnyAttr()
            if constraint is None
            else irdl_to_attr_constraint(constraint, allow_type_var=True),
            converter,
        ),
    )

def check_attr_name(cls: type):
    """Check that the attribute class has a correct name."""
    name = None
    for base in cls.mro():
        if "name" in base.__dict__:
            name = base.__dict__["name"]
            break

    if not isinstance(name, str):
        raise PyRDLAttrDefinitionError(
            f"pyrdl attribute definition '{cls.__name__}' does not "
            "define the attribute name. The attribute name is defined by "
            "adding a 'name' field with a string value."
        )

    dialect_attr_name = name.split(".")
    if len(dialect_attr_name) >= 2:
        return

    if not issubclass(cls, BuiltinAttribute):
        raise PyRDLAttrDefinitionError(
            f"Name '{name}' is not a valid attribute name. It should be of the form "
            "'<dialect>.<name>'."
        )

def irdl_param_attr_definition(cls: _PAttrTT) -> _PAttrTT:
    """Decorator used on classes to define a new attribute definition."""

    attr_def = ParamAttrDef.from_pyrdl(cls)

    @classmethod
    def get_irdl_definition(cls: type[ParametrizedAttribute]):
        return attr_def

    cls.get_irdl_definition = get_irdl_definition

    if issubclass(cls, TypedAttribute):
        type_indexes = tuple(
            i for i, (p, _) in enumerate(attr_def.parameters) if p == "type"
        )
        if not type_indexes:
            raise PyRDLAttrDefinitionError(
                f"TypedAttribute {cls.__name__} should have a 'type' parameter."
            )
        type_index = type_indexes[0]

        @classmethod
        def get_type_index(cls: Any) -> int:
            return type_index

        setattr(cls, "get_type_index", get_type_index)

    return runtime_final(dataclass(frozen=True, init=False)(cls))

@dataclass_transform(frozen_default=True, field_specifiers=(param_def,))
def irdl_attr_definition(
    cls: TypeAttributeInvT | None = None, *, init: bool = True
) -> TypeAttributeInvT | Callable[[TypeAttributeInvT], TypeAttributeInvT]:
    def decorator(cls: TypeAttributeInvT) -> TypeAttributeInvT:
        check_attr_name(cls)
        if issubclass(cls, ParametrizedAttribute):
            return irdl_param_attr_definition(cls)
        if issubclass(cls, Data):
            # This used to be convoluted
            # But Data is already frozen itself, so any child Attribute still throws on
            # .data!
            return runtime_final(cast(TypeAttributeInvT, cls))
        raise TypeError(
            f"Class {cls.__name__} should either be a subclass of 'Data' or "
            "'ParametrizedAttribute'"
        )

    if cls is None:
        return decorator
    return decorator(cls)

def irdl_list_to_attr_constraint(
    pyrdl_constraints: Sequence[IRDLAttrConstraint],
    *,
    allow_type_var: bool = False,
) -> AttrConstraint:
    """
    Convert a list of PyRDL type annotations to an AttrConstraint.
    Each list element correspond to a constraint to satisfy.
    If there is a `ConstraintVar` annotation, we add the entire constraint to
    the constraint variable.
    """
    constraints = tuple(
        irdl_to_attr_constraint(arg, allow_type_var=allow_type_var)
        for arg in pyrdl_constraints
    )

    if len(constraints) > 1:
        return AllOf(constraints)

    if not constraints:
        return AnyAttr()

    return constraints[0]

def irdl_to_attr_constraint(
    irdl: IRDLAttrConstraint[AttributeInvT],
    *,
    allow_type_var: bool = False,
) -> AttrConstraint[AttributeInvT]:
    if isinstance(irdl, AttrConstraint):
        return cast(AttrConstraint[AttributeInvT], irdl)

    if isinstance(irdl, Attribute):
        return cast(AttrConstraint[AttributeInvT], EqAttrConstraint(irdl))

    if isinstance(irdl, ConstraintConvertible):
        value = cast(ConstraintConvertible[AttributeInvT], irdl)
        return value.constr()

    # Annotated case
    # Each argument of the Annotated type corresponds to a constraint to satisfy.
    # If there is a `ConstraintVar` annotation, we add the entire constraint to
    # the constraint variable.
    if get_origin(irdl) == Annotated:
        return cast(
            AttrConstraint[AttributeInvT],
            irdl_list_to_attr_constraint(
                get_args(irdl),
                allow_type_var=allow_type_var,
            ),
        )

    # Attribute class case
    # This is an `AnyAttr`, which does not constrain the attribute.
    if irdl is Attribute:
        return cast(AttrConstraint[AttributeInvT], AnyAttr())

    # Attribute class case
    # This is a coercion for an `BaseAttr`.
    if (
        isclass(irdl)
        and not isinstance(irdl, GenericAlias)
        and issubclass(irdl, Attribute)
    ):
        return cast(AttrConstraint[AttributeInvT], BaseAttr(irdl))

    # Type variable case
    # We take the type variable bound constraint.
    if isinstance(irdl, TypeVar):
        if not allow_type_var:
            raise PyRDLTypeError("TypeVar in unexpected context.")
        if irdl.__bound__ is None:
            raise PyRDLTypeError(
                "Type variables used in IRDL are expected to be bound."
            )
        # We do not allow nested type variables.
        constraint = irdl_to_attr_constraint(irdl.__bound__)
        return cast(AttrConstraint[AttributeInvT], TypeVarConstraint(irdl, constraint))

    origin = get_origin(irdl)

    # Generic case
    if isclass(origin) and (
        issubclass(origin, GenericData)
        or (issubclass(origin, ParametrizedAttribute) and issubclass(origin, Generic))
    ):
        if issubclass(origin, GenericData):
            base_constr = origin.constr()
        else:
            base_constr = ParamAttrConstraint(
                origin,
                tuple(
                    param.constr for _, param in origin.get_irdl_definition().parameters
                ),
            )

        type_vars = get_type_var_from_generic_class(cast(type, origin))
        args: tuple[IRDLAttrConstraint | int | TypeForm[int], ...] = get_args(irdl)

        # Here, we use `__default__` to check for default values, instead of `has_default`.
        # This is so users that are not using typing-extensions can still use it, as well
        # as marimo notebooks, which seems to replace `typing-extensions` with `typing`.
        if len(args) < len(type_vars) and all(
            hasattr(arg, "__default__") for arg in type_vars[len(args) :]
        ):
            # Check for default values
            args += tuple(arg.__default__ for arg in type_vars[len(args) :])

        # Check that we have the right number of parameters
        if len(args) != len(type_vars):
            raise PyRDLTypeError(
                f"{origin.name} expects {len(type_vars)} parameters, got {len(args)}."
            )

        type_var_mapping = {
            type_var: get_constraint(arg, allow_type_var=allow_type_var)
            for type_var, arg in zip(type_vars, args)
        }

        return cast(
            AttrConstraint[AttributeInvT],
            base_constr.mapping_type_vars(type_var_mapping),
        )

    # Union case
    # This is a coercion for an `AnyOf` constraint.
    if origin == UnionType or origin == Union:
        constraints: list[AttrConstraint] = []
        for arg in get_args(irdl):
            constraints.append(
                irdl_to_attr_constraint(
                    arg,
                    allow_type_var=allow_type_var,
                )
            )
        return cast(AttrConstraint[AttributeInvT], AnyOf.get(*constraints))

    if isclass(irdl) and issubclass(irdl, ConstraintConvertible):
        attr_data = cast(type[ConstraintConvertible[AttributeInvT]], irdl)
        return attr_data.base_constr()

    if origin is Literal:
        literal_args = get_args(irdl)
        if any(isinstance(value, bool) for value in literal_args):
            raise PyRDLTypeError("Boolean literals are currently not supported")
        return AnyOf.get(*literal_args)

    # Better error messages for missing GenericData in Data definitions
    if isclass(origin) and issubclass(origin, Data):
        raise PyRDLTypeError(
            f"Generic `Data` type '{origin.name}' cannot be converted to "
            "an attribute constraint. Consider making it inherit from "
            "`GenericData` instead of `Data`."
        )

    raise PyRDLTypeError(f"Unexpected irdl constraint: {irdl}")

@deprecated("Use constraint `.get` constructors or irdl_to_attr_constraint directly")
def base(irdl: type[AttributeInvT]) -> AttrConstraint[AttributeInvT]:
    """
    Converts an attribute type into the equivalent constraint, detecting generic
    parameters if present.
    """
    return irdl_to_attr_constraint(irdl)

@deprecated("Use constraint `.get` constructors or irdl_to_attr_constraint directly")
def eq(irdl: AttributeInvT) -> AttrConstraint[AttributeInvT]:
    """
    Converts an attribute instance into the equivalent constraint.
    """
    return irdl_to_attr_constraint(irdl)

def get_optional_int_constraint(arg: Any) -> IntConstraint | None:
    """
    If the input is an int or an int type, return the corresponding constraint,
    otherwise return None.
    Boolean literals are currently unsupported and raise `PyRDLTypeError`.
    """
    if isinstance(arg, bool):
        raise PyRDLTypeError("Boolean literals are currently not supported")

    if isinstance(arg, int):
        return EqIntConstraint(arg)

    if isclass(arg) and issubclass(arg, int):
        return AnyInt()

    if get_origin(arg) is Literal:
        literal_args = get_args(arg)
        if any(isinstance(value, bool) for value in literal_args):
            raise PyRDLTypeError("Boolean literals are currently not supported")

        if all(isinstance(literal_arg, int) for literal_arg in literal_args):
            return IntSetConstraint.get(*literal_args)

    if get_origin(arg) in (Union, UnionType):
        # Check every member before returning None for a non-integer constraint.
        constraints = [get_optional_int_constraint(member) for member in get_args(arg)]
        ints: set[int] = set()
        for constraint in constraints:
            match constraint:
                case EqIntConstraint():
                    ints.add(constraint.value)
                case IntSetConstraint():
                    ints.update(constraint.values)
                case _:
                    return None
        return IntSetConstraint.get(*ints)

    if (
        isinstance(arg, TypeVar)
        and (base := get_optional_int_constraint(arg.__bound__)) is not None
    ):
        return IntTypeVarConstraint(arg, base)

def get_int_constraint(arg: "int | TypeForm[int]") -> IntConstraint:
    """
    If the input is an int or an int type, return the corresponding constraint,
    otherwise raise `PyRDLTypeError`.
    """
    if (ic := get_optional_int_constraint(arg)) is not None:
        return ic

    raise PyRDLTypeError(f"Unexpected int type: {arg}")

def range_constr_coercion(
    attr: (IRDLAttrConstraint[AttributeCovT] | RangeConstraint[AttributeCovT]),
) -> RangeConstraint[AttributeCovT]:
    if isinstance(attr, RangeConstraint):
        return cast(RangeConstraint[AttributeCovT], attr)
    return RangeOf.get(attr)

def single_range_constr_coercion(
    attr: IRDLAttrConstraint[AttributeCovT],
) -> RangeConstraint[AttributeCovT]:
    return RangeOf.get(attr).of_length(1)

def get_constraint(
    arg: "IRDLAttrConstraint | int | TypeForm[int]",
    *,
    allow_type_var: bool = False,
) -> AttrConstraint | IntConstraint:
    """
    Converts the input expression, constraint, or type to the corresponding constraint.
    """
    if (ic := get_optional_int_constraint(arg)) is not None:
        return ic

    return irdl_to_attr_constraint(arg, allow_type_var=allow_type_var)  # pyright: ignore[reportArgumentType]