core: Add shorthand type for types that can be converted to IRDL (#3836)

Stacked PRs:
 * __->__#3836
 * #3835


--- --- ---

### core: Add shorthand type for types that can be converted to IRDL


There was a lot of inconsistencies in the typing of functions that are
using the type to constraint conversion. Having a shorthand makes this
more explicit.

xdsl/dialects/builtin.py

@@ -1818,7 +1818,7 @@ def print(self, printer: Printer) -> None:
     "_MemRefTypeElement", bound=Attribute, covariant=True, default=Attribute
 )
 _UnrankedMemRefTypeElems = TypeVar(
-    "_UnrankedMemRefTypeElems", bound=Attribute, covariant=True
+    "_UnrankedMemRefTypeElems", bound=Attribute, covariant=True, default=Attribute
 )
 _UnrankedMemRefTypeElemsInit = TypeVar("_UnrankedMemRefTypeElemsInit", bound=Attribute)
 

xdsl/dialects/mpi.py

@@ -3,7 +3,9 @@
 from abc import ABC
 from collections.abc import Sequence
 from enum import Enum
-from typing import Generic, TypeVar
+from typing import Generic
+
+from typing_extensions import TypeVar
 
 from xdsl.dialects import llvm
 from xdsl.dialects.builtin import (
@@ -110,7 +112,7 @@ class DataType(ParametrizedAttribute, TypeAttribute):
 
 VectorWrappable = RequestType | StatusType | DataType
 VectorWrappableConstr = base(RequestType) | base(StatusType) | base(DataType)
-_VectorT = TypeVar("_VectorT", bound=VectorWrappable)
+_VectorT = TypeVar("_VectorT", bound=VectorWrappable, default=VectorWrappable)
 
 
 @irdl_attr_definition

xdsl/dialects/stencil.py

@@ -5,7 +5,9 @@
 from itertools import pairwise
 from math import prod
 from operator import add, lt, neg
-from typing import Annotated, Generic, TypeAlias, TypeVar, cast
+from typing import Annotated, Generic, TypeAlias, cast
+
+from typing_extensions import TypeVar
 
 from xdsl.dialects import builtin, memref
 from xdsl.dialects.builtin import (
@@ -73,7 +75,9 @@
 from xdsl.utils.hints import isa
 from xdsl.utils.isattr import isattr
 
-_FieldTypeElement = TypeVar("_FieldTypeElement", bound=Attribute, covariant=True)
+_FieldTypeElement = TypeVar(
+    "_FieldTypeElement", bound=Attribute, covariant=True, default=Attribute
+)
 
 
 @irdl_attr_definition

xdsl/interpreter.py

@@ -251,11 +251,14 @@ def annot(
     return annot
 
 
+AttributeInvNoDefaultT = TypeVar("AttributeInvNoDefaultT", bound=Attribute)
+
+
 def impl_attr(
-    input_type: type[AttributeInvT],
+    input_type: type[AttributeInvNoDefaultT],
 ) -> Callable[
-    [AttrImpl[_FT, AttributeInvT]],
-    AttrImpl[_FT, AttributeInvT],
+    [AttrImpl[_FT, AttributeInvNoDefaultT]],
+    AttrImpl[_FT, AttributeInvNoDefaultT],
 ]:
     """
     Marks the conversion from an attribute to a Python value. The

xdsl/ir/core.py

@@ -21,14 +21,13 @@
     Generic,
     NoReturn,
     Protocol,
-    TypeVar,
     cast,
     get_args,
     get_origin,
     overload,
 )
 
-from typing_extensions import Self, deprecated
+from typing_extensions import Self, TypeVar, deprecated
 
 from xdsl.traits import IsTerminator, NoTerminator, OpTrait, OpTraitInvT
 from xdsl.utils.exceptions import VerifyException
@@ -338,8 +337,10 @@ class SpacedOpaqueSyntaxAttribute(OpaqueSyntaxAttribute):
 
 DataElement = TypeVar("DataElement", covariant=True, bound=Hashable)
 
-AttributeCovT = TypeVar("AttributeCovT", bound=Attribute, covariant=True)
-AttributeInvT = TypeVar("AttributeInvT", bound=Attribute)
+AttributeCovT = TypeVar(
+    "AttributeCovT", bound=Attribute, covariant=True, default=Attribute
+)
+AttributeInvT = TypeVar("AttributeInvT", bound=Attribute, default=Attribute)
 
 
 @dataclass(frozen=True)

xdsl/irdl/attributes.py

@@ -11,9 +11,11 @@
 from inspect import isclass
 from types import FunctionType, GenericAlias, UnionType
 from typing import (
+    TYPE_CHECKING,
     Annotated,
     Any,
     Generic,
+    TypeAlias,
     TypeVar,
     Union,
     cast,
@@ -22,6 +24,10 @@
     get_type_hints,
 )
 
+if TYPE_CHECKING:
+    from typing_extensions import TypeForm
+
+
 from xdsl.ir import (
     Attribute,
     AttributeInvT,
@@ -256,8 +262,30 @@ def irdl_attr_definition(cls: TypeAttributeInvT) -> TypeAttributeInvT:
     )
 
 
+IRDLGenericAttrConstraint: TypeAlias = (
+    GenericAttrConstraint[AttributeInvT]
+    | Attribute
+    | type[AttributeInvT]
+    | "TypeForm[AttributeInvT]"
+    | ConstraintVar
+    | TypeVar
+)
+"""
+Attribute constraints represented using the IRDL python frontend. Attribute constraints
+can either be:
+- An instance of `AttrConstraint` representing a constraint on an attribute.
+- An instance of `Attribute` representing an equality constraint on an attribute.
+- A type representing a specific attribute class.
+- A TypeForm that can represent both unions and generic attributes.
+- A `ConstraintVar` representing a constraint variable.
+"""
+
+IRDLAttrConstraint = IRDLGenericAttrConstraint[Attribute]
+"""See `IRDLGenericAttrConstraint`."""
+
+
 def irdl_list_to_attr_constraint(
-    pyrdl_constraints: Sequence[Any],
+    pyrdl_constraints: Sequence[IRDLAttrConstraint],
     *,
     allow_type_var: bool = False,
     type_var_mapping: dict[TypeVar, AttrConstraint] | None = None,
@@ -299,7 +327,7 @@ def irdl_list_to_attr_constraint(
 
 
 def irdl_to_attr_constraint(
-    irdl: Any,
+    irdl: IRDLAttrConstraint,
     *,
     allow_type_var: bool = False,
     type_var_mapping: dict[TypeVar, AttrConstraint] | None = None,